AtCoder Grand Contest 014E: Blue and Red Tree

http://agc014.contest.atcoder.jp/tasks/agc014_e

マージテクの問題を一つ解きたくなったので。

詳しくは解説に任せるが、青と赤の辺が重なっている部分を一つの頂点に併合していく感じの操作が出来れば良い。
これは各頂点から生えてる辺を持ってマージテクをしていくことによって実現できる。

実装はendagorionさんのを参考にした。
自力で実装しようとしたら、頂点をマージテクしようとしてうまくできなかった。
グラフもいつものvector G[MAX_N]みたいな形で表現してしまい、mergeしたときの処理が出来なかった。
変に赤と青を区別してしまったため、シンプルな実装からより遠ざけられた。

setの使い方もよくわかっていなくて、イテレータがランダムアクセスできないことを知った。

UF U;
int N;

vector<pi> L[MAX_N];
pi E[2][MAX_N];
set<pi> S[MAX_N];

void solve() {
	cin >> N;
	U.init(N);
	rep(k, 0, 2) {
		rep(i, 0, N - 1) {
			int a, b;
			cin >> a >> b; a--; b--;
			if(a > b) swap(a, b);
			E[k][i] = pi(a, b);
			S[k].insert(pi(a, b));
			L[a].push_back(pi(i, k));
			L[b].push_back(pi(i, k));
		}
	}

	queue<pi> que;
	for(auto w: S[0]) {
		if(S[1].count(w)) que.push(w);
	}

	while(!que.empty()) {
		pi p = que.front(); que.pop();
		int a = U.find(p.fst), b = U.find(p.sec);
		if(a == b) continue;
		if((int)L[a].size() > (int)L[b].size()) swap(a, b);
		U.unite(b, a);
		for(auto w: L[a]) {
			int id = w.fst, k = w.sec;
			S[k].erase(E[k][id]);
			E[k][id].fst = U.find(E[k][id].fst);
			E[k][id].sec = U.find(E[k][id].sec);
			if(E[k][id].fst > E[k][id].sec) swap(E[k][id].fst, E[k][id].sec);
			S[k].insert(E[k][id]);
			if(S[1 - k].count(E[k][id])) que.push(E[k][id]);
			L[b].push_back(w);
		}
		L[a].clear();
	}

	bool ok = true;
	for(int i = 0; i < N; i++) {
		if(U.find(0) != U.find(i)) ok = false;
	}
	cout << (ok ? "YES" : "NO") << "\n"; 
}

この実装ほんとに頭が良すぎる。