AtCoder ABC #334 (ユニークビジョンプログラミングコンテスト2023 クリスマス) : G - Christmas Color Grid 2

650ptの割にすんなり。
https://atcoder.jp/contests/abc334/tasks/abc334_g

問題

赤緑で2色に塗られたグリッドが与えられる。
緑セルを等確率で1つ赤にした時、緑セルの連結成分の個数の期待値を求めよ。

解法

巻き戻し可能なUnionFindを使い、入力の状態から緑セルを1つ除いた状態を総当たりする。
緑セルがG個あるとする。０～(G-1)の時間経過が起こるとして、時刻hにはある緑セル1つが赤くなっている状態になるとする。
各時刻で赤くなるセルは異なるようにすれば、各時刻の緑セルの連結成分を求められれば良い。

これをSegTreeの要領で処理する。
時刻hで緑セルが赤くなるということは、逆に[0,h-1]と[h+1,G-1]の時間ではそのセルは緑である。
SegTreeの要領で時間をノードに載せた二分木を考え、あるSubTree内でそのセルが緑であるようにする。
この二分木上をDFSしながら、緑セルを追加・削除していき、葉のノードに到達したらその時点の連結成分数を数えると良い。

int H,W;
string S[1010];
int T[1010][1010];
const ll mo=998244353;
int NG;

template<int um> class UF_pop {
	public:
	vector<int> par,rank; vector<vector<int>> hist;
	UF_pop() {par=rank=vector<int>(um,0); for(int i=0;i<um;i++) par[i]=i;}
	void reinit() {int i; hist.clear(); FOR(i,um) rank[i]=0,par[i]=i;}
	int operator[](int x) {return (par[x]==x)?(x):operator[](par[x]);}
	void pop() {
		if(hist.size()) {
			auto a=hist.back();
			hist.pop_back();
			par[a[0]]=a[1]; rank[a[0]]=a[2];
			par[a[3]]=a[4]; rank[a[3]]=a[5];
		}
	}
	void operator()(int x,int y) {
		x=operator[](x); y=operator[](y);
		hist.push_back({x,par[x],rank[x],y,par[y],rank[y]});
		if(x==y) return;
		if(rank[x]<rank[y]) par[x]=y;
		else rank[x]+=(rank[x]==rank[y]), par[y]=x;
	}
};
UF_pop<1<<22> uf;

vector<pair<int,int>> cand[1<<22];
ll sum;

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

void add(int cur,int TL,int TR,pair<int,int> p,int CL,int CR) {
	if(CL>=CR) return;
	if(TL>CR) return;
	if(TR<=CL) return;
	TL=max(TL,CL);
	TR=min(TR,CR);
	if(TL==CL&&TR==CR) {
		cand[cur].push_back(p);
	}
	else {
		if(CR-CL>1) {
			int CM=(CL+CR)/2;
			add(cur*2,TL,TR,p,CL,CM);
			add(cur*2+1,TL,TR,p,CM,CR);
		}
	}
}

void dfs(int cur,int CL,int CR,int num) {
	if(CL>=CR) return;
	int step=0;
	int d[4]={0,1,0,-1},i;
	int cnum=num;
	FORR2(cy,cx,cand[cur]) {
		T[cy][cx]=1;
		cnum++;
		FOR(i,4) {
			int ty=cy+d[i];
			int tx=cx+d[i^1];
			if(T[ty][tx]&&uf[ty*2000+tx]!=uf[cy*2000+cx]) {
				step++;
				cnum--;
				uf(ty*2000+tx,cy*2000+cx);
			}
		}
	}
	if(CL+1==CR) {
		sum+=cnum;
	}
	else {
		int CM=(CL+CR)/2;
		dfs(cur*2,CL,CM,cnum);
		dfs(cur*2+1,CM,CR,cnum);
	}
	FORR2(cy,cx,cand[cur]) {
		T[cy][cx]=0;
	}
	while(step--) uf.pop();
}


void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>H>>W;
	int NR=0;
	FOR(y,H) {
		cin>>S[y+1];
		FORR(c,S[y+1]) if(c=='#') NG++;
		S[y+1]='.'+S[y+1]+'.';
	}
	S[0]=S[H+1]=string(W+2,'.');
	H+=2;
	W+=2;
	
	int id=0;
	FOR(y,H) FOR(x,W) if(S[y][x]=='#') {
		add(1,0,id,{y,x},0,NG);
		add(1,id+1,NG,{y,x},0,NG);
		id++;
	}
	
	
	dfs(1,0,NG,0);
	
	
	
	sum=sum%mo*modpow(NG);
	cout<<sum%mo<<endl;
	
}

まとめ

このテク最近Codeforcesで知ったので、ちゃんと使えて良かった。