[AtCoder] ABC 157 D – Friend Suggestions

2020年3月4日2020年12月15日

問題

方針

まず初めに、友達関係の Union-Find を作成します。このとき、人 \( i \) の友達の数を \( f_i \) として数えます。次に、人 \( i \) のブロック人数を \( b_i \) と、人 \( i \) と 人 \( j \) がブロック関係かつ 人 \( i \) と 人 \( j \) が同じグラフに所属してるとき、人 \( i \) のブロックの人数を \( b_i \) として加算します。また、\( b_j \) も加算します。人 \( i \) の連結成分の数を \( a_i \) とすると、人 \( i \) の友達候補は、\( a_i – f_i – b_i – 1 \) となります。

コード

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

// Union-Find
class DisjointSet {
public:
  vector<int> p, rank, size;
  int group_num;
  DisjointSet(int n) {
    p.assign(n, 0);
    rank.assign(n, 0);
    size.assign(n, 0);
    group_num = n;
    for (int i = 0; i < n; i++) {
      make_set(i);
    }
  }

  void make_set(int x) {
    p[x] = x;
    rank[x] = 0;
    size[x] = 1;
  }

  bool same(int x, int y) {
    return find_set(x) == find_set(y);;
  }

  int find_set(int x) {
    if (x != p[x]) {
      p[x] = find_set(p[x]);
    }
    return p[x];
  }

  int union_size(int x) {
    if (x != p[x]) {
      size[x] = union_size(find_set(x));
    }
    return size[x];
  }

  void link(int x, int y) {
    if(rank[x]  > rank[y]) {
      p[y] = x;
      size[x] = size[x] + size[y];
    }else if(rank[x] < rank[y]){
      p[x] = y;
      size[y] = size[x] + size[y];
    }else if(x != y) {
      p[y] = x;
      rank[x]++;
      size[x] = size[x] + size[y];
    }
  }

  void unite(int x, int y) {
    if (find_set(x) != find_set(y)) group_num--;
    link(find_set(x), find_set(y));
  }
};

int main() {
    int N, M, K;
    cin >> N >> M >> K;
    int f[N]{};
    int b[N]{};
    DisjointSet ds(N);
    int A, B, C, D;
    for (int i = 0; i < M; i++) {
        cin >> A >> B;
        A--;
        B--;
        ds.unite(A, B);
        f[A]++;
        f[B]++;
    }
    for (int i = 0; i < K; i++) {
        cin >> C >> D;
        C--;
        D--;
        if (ds.same(C, D)) {
            b[C]++;
            b[D]++;
        }
    }

    for (int i = 0; i < N; i++) {
        int p = ds.union_size(i) - f[i] - b[i] - 1;
        if (i == N - 1) {
            cout << p << "\n";
        } else {
            cout << p << " ";
        }
    }
    return 0;
}