Codeforces ECR #121 : E. Black and White Tree

割と難しめな回だったみたいだけど、Eはすんなり解けているな…。
https://codeforces.com/contest/1626/problem/E

問題

木を成す無向グラフが与えられる。
各点は、白黒のいずれかで塗られている。

今点xに駒があるとき、黒い点yを1つ選び、xからyに向けて辺1つ分駒を移動させることを考える。
上記処理を任意回数行えるが、同じ点yを2回連続では選べないものとする。

駒を、どこか任意の黒点まで移動させることが可能かどうか、全点をそれぞれ初期位置とした場合について答えよ。

解法

各点を根としてみると、各隣接点のSubTreeに黒点が2個以上あれば、任意のタイミングで駒をそちらに遷移できる。
そう見なして、元のグラフから任意に遷移可能な有向辺を逆向きに残したグラフを考える。
この逆向きの辺でたどれる点は、黒点に遷移可能な位置を示すので、DFSでたどっていけばよい。

int N;
int C[303030],SC[303030],TC,can[303030];
vector<int> E[303030];
vector<int> OE[303030],RE[303030];
int vis[303030];

int dfs(int cur, int pre) {
	SC[cur]=C[cur];
	FORR(e,E[cur]) if(e!=pre) {
		int x=dfs(e,cur);
		if(x>=2) {
			RE[e].push_back(cur);
		}
		SC[cur]+=x;
	}
	FORR(e,E[cur]) if(e==pre&&TC-SC[cur]>=2) {
		RE[e].push_back(cur);
	}
	return SC[cur];
}

void dfs2(int cur) {
	if(vis[cur]) return;
	can[cur]=vis[cur]=1;
	FORR(e,RE[cur]) dfs2(e);
}

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	FOR(i,N) {
		cin>>C[i];
		TC+=C[i];
		can[i]=C[i];
	}
	FOR(i,N-1) {
		cin>>x>>y;
		x--,y--;
		E[x].push_back(y);
		E[y].push_back(x);
		if(C[x]) can[y]=1;
		if(C[y]) can[x]=1;
	}
	dfs(0,0);
	FOR(i,N) if(can[i]) dfs2(i);
	
	FOR(i,N) cout<<can[i]<<" ";
	cout<<endl;
	
}

まとめ

本番よくこれすんなり通せたな。