AtCoder Beginner Contest 049: D - 連結

問題

問題文

https://atcoder.jp/contests/abc049/tasks/arc065_b

問題概要

頂点数が n の無向グラフが与えられる. 2つの種類の辺があり, それぞれ M本の辺, L本の辺がある.

1種類のみの辺を使ったとき, 各頂点からいずれの種類の辺のみを使ってもたどり着ける頂点の数を答えよ.

解答例

指針

UnionFind
連想配列による集計

解説

それぞれの種類の辺について, UnionFind 木を用いれば, 連結性を求めることは容易にできる. それぞれ, uf1, uf2 というUnionFind木で管理することにする.

以下のように各頂点から, それぞれ別の頂点と連結かどうかを調べれば答えは求めることができるが, この場合の計算量は $O(n^{2} \log{(n)})$ となり間に合わない.

TLE解

// Time Limit Exceeded
// submission: https://atcoder.jp/contests/abc049/submissions/11023252
int main() {
  // (略)    
  for (int i = 0; i < n; i++) {
    int ans = 0;
    for (int j = 0; j < n; j++) {
      if (uf1.same(i, j) && uf2.same(i, j)) {
        ans++;
      }
    }
    cout << ans;
    cout << (i == n-1 ? "\n" : " ");
  }
  return 0
}

以下 UnionFind において, 頂点の根を返すメンバ関数の関数名を root, 同じ集合に属するかを判定するメンバ関数の関数名を same であるとして説明する.

ここで頂点 $i$ と uf1, uf2 のいずれにおいても連結となる頂点 $j$ とは, uf1.same(i, j) && uf2.same(i, j) が true となるものであるが, これは, uf1.root(i) == uf1.root(j) && uf2.root(i) == uf2.root(j) が true となるもの (根/代表元が等しい) であり, C++のmapなどの連想配列により, uf1.root(i) と uf2.root(i) をペアにして集計し, 同じペアのものを数えれば, uf1.root(i) == uf1.root(j) && uf2.root(i) == uf2.root(j) となる頂点 $j$ の個数を数える事ができる.

このときの時間計算量は $O(n\log{n})$ となる.

実装例

C++ ( https://atcoder.jp/contests/arc065/submissions/11023532 )

#include <iostream>
#include <algorithm>
#include <vector>
#include <map>

using namespace std;

struct UnionFind {
  std::vector<int> parent;
  int __size;
  UnionFind(int size_) : parent(size_, -1), __size(size_) {}
  void unite(int x, int y) {
    if ((x = root(x)) != (y = root(y))) {
      if (parent[y] < parent[x]) std::swap(x, y);
      parent[x] += parent[y];
      parent[y] = x;
      __size--;
    }
  }
  bool same(int x, int y) { return root(x) == root(y); }
  int root(int x) { return parent[x] < 0 ? x : parent[x] = root(parent[x]); }
  int size(int x) { return -parent[root(x)]; }
  int size() { return __size; }
};

int main() {
  int n, k, l;
  cin >> n >> k >> l;

  UnionFind uf1(n);
  for (int i = 0; i < k; i++) {
    int a, b; cin >> a >> b;
    uf1.unite(a-1, b-1);
  }

  UnionFind uf2(n);
  for (int i = 0; i < l; i++) {
    int a, b; cin >> a >> b;
    uf2.unite(a-1, b-1);
  }

  map<pair<int, int>, int> mp;
  for (int i = 0; i < n; i++) {
    mp[make_pair(uf1.root(i), uf2.root(i))]++;
  }

  for (int i = 0; i < n; i++) {
    cout << mp[make_pair(uf1.root(i), uf2.root(i))];
    cout << (i == n-1 ? "\n" : " ");
  }

  return 0;
}