Codeforces #947 : D. Paint the Tree

まぁまぁの出来だった回。
https://codeforces.com/contest/1975/problem/D

問題

N頂点の木を成すグラフがあり、初期状態で全点白く塗られている。
うち2つの点に2つのチェスの駒があり、前者のいる頂点は赤、後者のいる頂点は青に塗られている。
ここで以下を交互に繰り返す。

• 前者の駒を隣接点に動かす。移動先の頂点が白なら赤くする
• 後者の駒を隣接点に動かす。移動先の頂点が赤なら青くする

全点を青にするまでにかかる最小ステップを求めよ。

解法

両者が最速で合流させたうえで、そこからの移動回数を考える。
合流後、全点たどって元の位置に戻ると両駒2*(N-1)回移動が必要だが、最後元の位置に戻る必要がないので、最後を最遠点を塗りつぶすことにすると、2*(N-1)から最遠点までの距離を引けばよい。

int T,N;
int A,B;
vector<int> E[202020];
int D[3][202020];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>T;
	while(T--) {
		cin>>N;
		cin>>A>>B;
		A--,B--;
		FOR(i,N) {
			E[i].clear();
			D[0][i]=D[1][i]=D[2][i]=1<<20;
		}
		FOR(i,N-1) {
			cin>>x>>y;
			E[x-1].push_back(y-1);
			E[y-1].push_back(x-1);
		}
		D[0][A]=D[1][B]=0;
		queue<int> Q;
		Q.push(A);
		Q.push(N+B);
		while(Q.size()) {
			int id=Q.front()/N;
			int cur=Q.front()%N;
			Q.pop();
			FORR(e,E[cur]) if(chmin(D[id][e],D[id][cur]+1)) Q.push(id*N+e);
		}
		int ret=(D[0][B]+1)/2;
		FOR(i,N) {
			if((D[0][i]+D[1][i]==D[0][B]&&D[0][i]==D[1][i])||(D[0][i]+D[1][i]==D[0][B]&&D[0][i]+1==D[1][i])) x=i;
		}
		FOR(i,N) D[0][i]=1<<20;
		ret+=2*(N-1);
		D[0][x]=0;
		Q.push(x);
		int ma=0;
		while(Q.size()) {
			int id=Q.front()/N;
			int cur=Q.front()%N;
			Q.pop();
			ma=max(ma,D[0][cur]);
			FORR(e,E[cur]) if(chmin(D[id][e],D[id][cur]+1)) Q.push(id*N+e);
		}
		ret-=ma;
		cout<<ret<<endl;
		
	}
}

まとめ

コードがちょっと冗長だな…。