ABC 049 D 連結 解説してみた

ABC 049 D 連結が解けたのが嬉しかったので解説してみた(ちなみに43分かかった)。

Union-Find木を使う問題である。 タイトルからして、「ああ、グラフの連結成分求める問題か。Union-Find使えば楽勝じゃん」とか思って取り掛かった。

・・・・・・甘かった。流石にD問題がそこまで単純な問題であるわけがなかった。 詳細はリンクを辿ってほしいのだが、この問題には２種類の「連結」がある。 「道路」による連結と「鉄道」による連結があり、これらは区別されなければならない。 しかし、最終的にこれらの両方で連結になっている相手を求めるという問題である。

連結かどうかの情報を保持するのにUnion-Find木を用いるのは分かる。 道路と鉄道、それぞれに対応する計２つのUnion-Find木を用意する。 ここで問題となるのは、いかにして「両方で連結」になっているかどうかの情報を効率よく抽出するかである。 勿論、最もナイーブなやり方は全てのノードの組み合わせについてループを回して、一つ一つ調べる方法である。 しかし、これではNが大きいときにTLEしてしまう。

Union-Find木の基本的な機能の一つに、「ノードの根を求める」というのがある。 ノードuとvが同じ木に属すということは、これらの根が同じであることと同値である。 このことを用いて、何とか都市をグループ分けできないだろうか。

道路による連結を表すUnion-Find木をuf1, 鉄道による連結を表すUnion-Find木をuf2とする。 都市uとvが道路と鉄道の両方で連結であることは、言い換えればこれらのuf1における根とuf2における根がそれぞれ等しいということである：

    uf1.root(u) == uf1.root(v) && uf2.root(u) == uf2.root(v)

そこで、第一の添字をuf1における根、第二の添字をuf2における根にした二次元配列を考え、配列の要素に都市の数を格納すればいい！！

...では駄目なのである。実際、上の指針に従ってn*nサイズのint型の二次元配列を定義するコードを投げると、サンプルは通るが途中のテストケースでREになってしまう。 理由は、メモリオーバーフローである。配列のサイズが大きすぎて領域が確保できないのだ。 ちなみに、これは配列をグローバル変数にしてヒープ領域に確保するようにしても同じである。ここには載せないが、コンパイル時に何かもう目が痛くなるような大量のエラーメッセージを見ることができる。

ではどうするか。ここは、配列ではなくmapを使うのがうまい方法だ。これなら、存在しないキーの組み合わせについて無駄なメモリを消費することがない。

    #include<cstdio>
    #include<bits/stdc++.h>
    typedef long long int ll;
    #define REP(i, n) for(int i = 0; i < (n); i++)
    #define FOR_IN(i, a, b) for(int i = (a); i < (b); i++)
    #define BETWEEN(x, a, b) ((x) >= (a) && (x) <= (b))
    #define BIT(b, i) (((b) >> (i)) & 1)
    #define LOG_F 1
    #define LOG(...) if(LOG_F) fprintf(stderr, __VA_ARGS__)
     
    using namespace std;
     
    ll pow(int x, int n){
      return n == 0 ? 1 : x * pow(x, n - 1);
    }
    struct union_find{
      vector<int> par;
      vector<int> rank;
      union_find(int n){
        rank = vector<int>(n,0);
        for(int i=0; i<n; i++){
          par.push_back(i);
        }
      }
     
      int root(int x){
        if(par[x]==x){
          return x;
        } else {
          return par[x] = root(par[x]);
        }
      }
     
      void unite(int x, int y){
        int rx = root(x);
        int ry = root(y);
     
        if(rx==ry)
          return;
     
        if(rank[rx]<rank[ry]){
          par[rx] = ry;
        } else {
          par[ry] = rx;
          if(rank[rx]==rank[ry])
            rank[rx]++;
        }
      }
     
      bool equiv(int x, int y){
        return root(x)==root(y);
      }
    };
     
     
    /*
    n k l
    p1 q1
    ...
    pk qk
    r1 s1
    ...
    rl sl
     
    n 2*10^5
    k.l 10^5
    p q r s <= n
    p < q
    r < s
     */
    int n,k,l;
     
    int main(){
      cin>>n>>k>>l;
      union_find uf1(n);
      union_find uf2(n);
      REP(i,k){
        int p,q;
        cin>>p>>q;
        p--;
        q--;
        uf1.unite(p,q);
      }
      REP(j,l){
        int r,s;
        cin>>r>>s;
        r--;
        s--;
        uf2.unite(r,s);
      }
      map<int, map<int, int>> root;
     
      REP(i,n)
        root[uf1.root(i)][uf2.root(i)]++;
     
      REP(i,n)
        cout << root[uf1.root(i)][uf2.root(i)] << " ";
     
      cout << endl;
      return 0;
    }