AtCoder Beginner Contest 049: D - 連結

問題

問題文

問題概要

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

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

解答例

指針

解説

それぞれの種類の辺について, 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 であるとして説明する.

ここで頂点 iuf1, 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}) となる.

実装例

#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;
}

感想

Problems みたら Difficulty 2003 となっていたけど, いま出題されたら 1200 程度と算出されそうだと思った.

「各UnionFind木における根のペアが等しいような頂点集合は, 各UnionFind木において, 互いに連結」 というのは, わざわざ言うまでもないくらい自明なことだけど気づけなかった.

参考文献