Atcoder Grand Contest 014 E. Blue and Red Tree

Atcoder Grand Contest 014 E. Blue and Red Tree

E: Blue and Red Tree - AtCoder Grand Contest 014 | AtCoder

感想

難しいがおもしろい。解説を読んで、少し人のコードを参考にして書いた。

解法

辺の色が異なる2つの木を重ね合せて考える。そのグラフ上で、ある2頂点$\ a, b\ $に対して青と赤の辺がどちらも張られているとする。この2頂点に関する操作はどのタイミングでも可能なので、あとの方でやった方がよい。

さらにこれらの2点間は、他の点での操作をするときに、いつでも移動可能であるため、この2点に辺を張る操作を「最後にする」という意味をもたせて、縮約して考えてみる。この新しくできたグラフの上でまた同様のことを考えていく。 もしこの縮約を繰り返していって一つの頂点にまとめてしまえるならば、その縮約をしていった逆の順番で、辺を構成していくことが可能であるといえる。

逆に頂点が2個以上残っているある段階でそれ以上縮約できない状態になったとしたら、小さい図を書いて実験してみると、なんとなくNOになりそうだとわかる。木である性質を使って色々場合分けしたら証明できると思う。ちゃんとした証明はeditorialで。

この縮約はある部分でやったからといって他の部分でできなくなるということはないので貪欲にやってよい。そのためにはいわゆるマージテクを使う。

簡単に辺を削除できるように、setで辺の情報をもつようにした。さらに、辺の色は、重複しているときのみ考えればよいため、色は実質無視して良い。

実装の方針としては、辺の重複が現れるたびに、queueなどのデータ構造にその頂点のペアを保管しておく。そこから取り出した頂点のペアが$\ (a, b)\ $だとすると、その頂点につながる辺の数が少ない方を、多い方にマージするようにする。多い方を$\ a\ $とすると、まず$\ a\ $と$\ b\ $の間の辺を取り除き、$\ b\ $とつながる各頂点を$\ a\ $とつなぎ、最後に$\ b\ $からの情報をすべてクリアする。

ただし、注意すべきことは、ある時点でpushされた頂点のペアが、popされるまでの間に縮約によって別の頂点に変化していることがある。それに対処するために、UnionFindのようなものを使う。すなわち、根が自身がマージされたノードであるような構造を持っておけば、矛盾なくマージされた情報がわかる。

細かいことだが、もう一つ着目するところがあって、途中にあるif (g[a].size() < g[b].size()) swap(a, b);は必要に思われるが、実はこれを書かなくても高速に動く（というかあまりswapが実行されない）。実はemplace(a, c)で追加されるときにまだ十分にマージされていないcの方が小さい確率が十分高く、その結果bothの中の(a, b)についてもbの方が小さい確率が十分に高いからである。

実装

#include "bits/stdc++.h"
using namespace std;

#define mp      make_pair

int par[101010];

int find(int x) { return par[x] == x ? x : par[x] = find(par[x]); }

int main() {
        int n;
        scanf("%d", &n);
        set<int> g[101010];
        queue<pair<int, int>> both;
        for (int i = 0; i < n; i ++) par[i] = i;
        for (int i = 0; i < 2 * n - 2; i ++) {
                int a, b;
                scanf("%d%d", &a, &b);
                a --, b --;
                if (g[a].count(b)) { 
                        both.emplace(a, b);
                } else { 
                        g[a].insert(b);
                        g[b].insert(a);
                }
        }
        int ans = 0;
        while (!both.empty()) {
                ans ++;
                int a, b;
                tie(a, b) = both.front();
                both.pop();
                a = find(a), b = find(b);
                //if (g[a].size() < g[b].size()) swap(a, b);
                a = par[a], b = par[b];
                if (a == b) continue;
                par[b] = a;
                g[a].erase(b);
                g[b].erase(a);
                for (auto c : g[b]) {
                        g[c].erase(b);
                        if (g[a].count(c)) { 
                                both.emplace(a, c);
                        } else { 
                                g[a].insert(c);
                                g[c].insert(a);
                        }
                }
                g[b].clear();
        }
        puts(ans == n - 1 ? "YES" : "NO");
        return 0;
}