けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

AtCoder ABC 157 D - Friend Suggestions (400 点)

とても教育的な Union-Find!!!

問題へのリンク

問題概要

 N 人がいて、 M 組の友達関係と、 K 組のブロック関係がある (いずれも双方向的)。

各人について、

  • 直接的な友達関係ではない
  • ブロック関係でもない
  • 友達の友達の...とたどっていくと到着できる

ような人の人数を求めよ。

制約

  •  N, M, K \le 10^{5}

考えたこと

まず友達関係を Union-Find によってグループ分けして考えることにしよう。つまり、人  i に対して、 i と同じ連結成分内に含まれる人のみに考察対象を絞って良いのだ。そうすると答えは

  • Union-Find 上で、人  i の属する連結成分に含まれる人数を  a とする
  • その連結成分内の人のうち、 i と友達であるか、ブロック関係にあるかのいずれかであるような人数を  b とする
  • このとき、答えは  a - b - 1 となる (-1 は、自分自身を除く分)

という風に求めることができる。 人  i の属する連結成分に含まれる人数  a は Union-Find で簡単に求めることができる。問題は  b だ。

i と同じ連結成分内の人のうち、i と友達またはブロック関係にある人数

これは、以下のようなデータ構造をもつことにした

  • dame[ i ] := i と同じ連結成分内の人であって、i と友達またはブロック関係にある人を格納した集合 (set 型で表現)

そして、友達関係 (a, b) やブロック関係 (c, d) を次のように処理することにした。

友達関係

  • dame[ a ].insert( b )
  • dame[ b ].insert( a )

ブロック関係

  • c と d が同じグループの場合に限り
    • dame[ c ].insert( d )
    • dame[ d ].insert( c )

 

こうしておけば、最終的な答えは

  • uf.size( i ) - dame[ i ].size() - 1 (Union-Find を uf とする)

で求めることができる。

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

struct UnionFind {
    vector<int> par;
    
    UnionFind(int n) : par(n, -1) { }

    int root(int x) {
        if (par[x] < 0) return x;
        else return par[x] = root(par[x]);
    }
    
    bool issame(int x, int y) {
        return root(x) == root(y);
    }
    
    bool merge(int x, int y) {
        x = root(x); y = root(y);
        if (x == y) return false;
        if (par[x] > par[y]) swap(x, y); // merge technique
        par[x] += par[y];
        par[y] = x;
        return true;
    }
    
    int size(int x) {
        return -par[root(x)];
    }
};

using pint = pair<int,int>;
int main () {
    int N, M, K;
    cin >> N >> M >> K;
    vector<set<int> > dame(N);
    UnionFind uf(N);
    for (int i = 0; i < M; ++i) {
        int a, b; cin >> a >> b; --a, --b;
        dame[a].insert(b);
        dame[b].insert(a);
        uf.merge(a, b);
    }
    for (int i = 0; i < K; ++i) {
        int c, d; cin >> c >> d; --c, --d;
        if (!uf.issame(c, d)) continue;
        dame[c].insert(d);
        dame[d].insert(c);
    }

    for (int i = 0; i < N; ++i) {
        int mem = uf.size(i) - 1; // 同じ連結成分の「自分以外」の人数
        mem -= dame[i].size(); // その中ですでに友達かブロック関係の人数
        cout << mem << " ";
    }
    cout << endl;
}