AtCoder ABC 157 D - Friend Suggestions (水色, 400 点)

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

問題へのリンク

問題概要

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

各人について、

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

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

制約

• 

考えたこと

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

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

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

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