[AtCoder] ABC157 D – Friend Suggestions 解説（400点）

2020年3月3日AtCoderグラフ,競プロ,連結成分,Union-Find

問題概要

N 人がいて、M 組が友達関係、K 組がブロック関係である。それぞれに対して以下を満たす「友達候補」が何人いるか答えよ。

  • まだ友達関係ではない
  • ブロック関係ではない
  • 友達関係を辿ってたどり着ける

制約

  • \(2 ≤ N ≤ 10^5\)
  • \(0 ≤ M ≤ 10^5\)
  • \(0 ≤ K ≤ 10^5\)
  • 友達関係とブロック関係は重複しない

考え方

友達候補というのは、友達関係を辿って到達できる人物です。この性質はグラフの構造に非常によく似ています。

人間を頂点・友達関係を辺としたグラフとして考えると、「友達候補」は以下の条件を満たす必要があります。

  • 間に辺がない
  • パスがある（同じ連結成分内である）
  • ブロック関係ではない

よって、「友達候補」の数は以下のように計算可能です。

友達候補の数 = 連結成分のサイズ-1 – 自身からの辺(友達関係)の数 – 連結成分内のブロック関係の数

1を引くのは、自身を「友達候補」から外すためです。

自身からの辺(友達関係)の数 は問題に与えられています。

連結成分ごとのサイズや、連結成分内のブロック関係の数は、深さ優先探索や Union-Find Tree などを用いて取得可能です。

解法（DFS）

  1. DFS を用いて、「自身がどの連結成分に含まれるか」「連結成分ごとのサイズ」をあらかじめ求める
  2. 「友達候補の数 = 連結成分のサイズ-1 – 自身からの辺(友達関係)の数 – 連結成分内のブロック関係の数」を各人間に対して計算する

C++での実装例

#include <bits/stdc++.h>
#define LOOP(n) for (int _i = 0; _i < (n); _i++)
#define REP(i, n) for (int i = 0; i < (n); ++i)
#define RREP(i, n) for (int i = (n); i >= 0; --i)
#define FOR(i, r, n) for (int i = (r); i < (n); ++i)
#define ALL(obj) begin(obj), end(obj)
using namespace std;

struct Edge {
    int to;
};
using Graph = vector<vector<Edge>>;

int N, M, K;

int id;
// 深さ優先探索
vector<int> cc_id;   // どの連結成分か
vector<int> cc_num;  // 連結成分のサイズ
void dfs(const Graph &G, int v, int &cnt) {
    cc_id[v] = id;
    cnt++;
    for (auto e : G[v]) {
        if (cc_id[e.to] == -1) {
            dfs(G, e.to, cnt);
        }
    }
}

int main() {
    cin >> N >> M >> K;
    Graph G(N);  // グラフとして考える
    vector<int> A(M), B(M);
    REP(i, M) {
        cin >> A.at(i) >> B.at(i);
        A[i]--, B[i]--;
        G[A[i]].push_back({B[i]});
        G[B[i]].push_back({A[i]});
    }
    vector<int> C(K), D(K);
    REP(i, K) {
        cin >> C.at(i) >> D.at(i);
        C[i]--, D[i]--;
    }

    // どの連結成分に属し、サイズがいくつなのかを求める
    cc_id.assign(N, -1);  // 初期化
    id = 0;
    REP(i, N) {
        if (cc_id[i] == -1) {
            int cnt = 0;
            dfs(G, i, cnt);
            cc_num.push_back(cnt);
            id++;
        }
    }

    vector<int> ans(N);
    REP(i, N) { ans[i] = cc_num[cc_id[i]] - 1 - (int)G[i].size(); }
    REP(i, K) {  // 同じ連結成分にあるブロック関係を引く
        if (cc_id[C[i]] == cc_id[D[i]]) {
            ans[C[i]]--;
            ans[D[i]]--;
        }
    }

    REP(i, (int)ans.size()) {
        cout << ans[i];
        if (i != (int)ans.size() - 1) cout << " ";
    }
    cout << endl;

    return 0;
}

解法（Union-Find Tree）

  1. Union-Find Tree を用いて、1つの連結成分を一つの集合(木) としてまとめる
  2. 「友達候補の数 = 連結成分のサイズ-1 – 自身からの辺(友達関係)の数 – 連結成分内のブロック関係の数」を各人間に対して計算する
Union-Find Tree を理解する！素集合系を扱うデータ構造

C++での実装例

Tree のサイズを求めたいので、Union-Find Tree の併合時の工夫として、サイズを利用したものを使うと良いです。

#include <bits/stdc++.h>
#define LOOP(n) for (int _i = 0; _i < (n); _i++)
#define REP(i, n) for (int i = 0; i < (n); ++i)
#define RREP(i, n) for (int i = (n); i >= 0; --i)
#define FOR(i, r, n) for (int i = (r); i < (n); ++i)
#define ALL(obj) begin(obj), end(obj)
using namespace std;

int N, M, K;

/* UnionFind：素集合系管理の構造体(union by size)
    isSame(x, y): x と y が同じ集合にいるか。 計算量はならし O(α(n))
    unite(x, y): x と y を同じ集合にする。計算量はならし O(α(n))
    treeSize(x): x を含む集合の要素数。
*/
struct UnionFind {
    vector<int> size, parents;
    UnionFind() {}
    UnionFind(int n) {  // make n trees.
        size.resize(n, 0);
        parents.resize(n, 0);
        for (int i = 0; i < n; i++) {
            makeTree(i);
        }
    }
    void makeTree(int x) {
        parents[x] = x;  // the parent of x is x
        size[x] = 1;
    }
    bool isSame(int x, int y) { return findRoot(x) == findRoot(y); }
    bool unite(int x, int y) {
        x = findRoot(x);
        y = findRoot(y);
        if (x == y) return false;
        if (size[x] > size[y]) {
            parents[y] = x;
            size[x] += size[y];
        } else {
            parents[x] = y;
            size[y] += size[x];
        }
        return true;
    }
    int findRoot(int x) {
        if (x != parents[x]) {
            parents[x] = findRoot(parents[x]);
        }
        return parents[x];
    }
    int treeSize(int x) { return size[findRoot(x)]; }
};

int main() {
    cin >> N >> M >> K;
    UnionFind uf(N);              // Union-Find Tree で考える
    vector<int> edges_num(N, 0);  // 自身から出る辺の数
    vector<int> A(M), B(M);
    REP(i, M) {
        cin >> A.at(i) >> B.at(i);
        A[i]--, B[i]--;
        edges_num[A[i]]++;
        edges_num[B[i]]++;
        uf.unite(A[i], B[i]);
    }
    vector<int> C(K), D(K);
    REP(i, K) {
        cin >> C.at(i) >> D.at(i);
        C[i]--, D[i]--;
    }

    vector<int> ans(N);
    REP(i, N) { ans[i] = uf.treeSize(i) - 1 - edges_num[i]; }
    REP(i, K) {  // 同じ連結成分にあるブロック関係を引く
        if (uf.isSame(C[i], D[i])) {
            ans[C[i]]--;
            ans[D[i]]--;
        }
    }

    REP(i, (int)ans.size()) {
        cout << ans[i];
        if (i != (int)ans.size() - 1) cout << " ";
    }
    cout << endl;

    return 0;
}

AtCoderグラフ,競プロ,連結成分,Union-Find