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

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

AtCoder ABC 016 C - 友達の友達 (試験管緑色)

グラフの入門に!

問題概要

頂点数  N、辺数  M の単純無向グラフが与えられます。

各頂点  v = 1, 2, \dots, N に対して、「友達の友達」、つまり頂点  v からの距離が 2 であるような頂点の個数を求めてください。

自分自身や、友達 (距離が 1 である頂点) は含めないものとします。

制約

  •  1 \le N \le 10
  •  0 \le M \le \frac{N(N-1)}{2}

解法

グラフのデータ構造を活用する練習問題です!

グラフのことが何も分からないという方は、次の記事を読んでみてください。

qiita.com

グラフをデータ構造として表す場合は、たとえば下図のように、各頂点に隣接する頂点のリストを考えます。

これらの情報は、二次元配列 vector<vector<int>> 型の変数 G を用いて管理できます。頂点  v に隣接する頂点の集合は G[v] と現せます。たとえば上図のグラフの場合、次のようになります。

G[0] = {5}
G[1] = {3, 6}
G[2] = {5, 7}
G[3] = {0, 7}
G[4] = {1, 2, 6}
G[5] = {}
G[6] = {7}
G[7] = {0}

今回の問題

このデータ構造を用いて、各頂点の「友達の友達」を求めていきましょう。

「友達の友達」から「友達」を除去する操作を効率よく実現するためには、集合型 (set 型) を用いるのが便利でしょう。

コード

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

int main() {
    int N, M;
    cin >> N >> M;
    vector<vector<int>> G(N);
    for (int i = 0; i < M; ++i) {
        int a, b;
        cin >> a >> b;
        --a, --b;  // 0-indexed にする
        G[a].push_back(b);
        G[b].push_back(a);
    }
    
    for (int v = 0; v < N; ++v) {
        set<int> S;
        
        // 頂点 v に隣接する頂点に隣接する頂点を見ていく
        for (auto i : G[v]) for (auto j : G[i]) S.insert(j);
        
        // 頂点 v 自身と、頂点 v に直接隣接する頂点は除く
        S.erase(v);
        for (auto i : G[v]) S.erase(i);
        
        cout << S.size() << endl;
    }
}