Pythonで友達リストからグループ数を求める方法｜DFSで連結成分を数えるアルゴリズム

友人関係のデータがリストとして与えられ、friends[i] には人 i の友達が誰なのかが格納されているとします。友情のつながりは双方向であり、各人は必ず自分自身とも友達であるものとみなします。また、互いの友人をたどる経路（パス）でつながっている2人は、同じ「友達グループ」に属すると定義されます。

このとき、全体にいくつの友達グループが存在するかを求めるのがこの問題です。

入力例と出力

たとえば次のような入力を考えます。

friends = [[0, 1, 5], [1, 0], [2], [3, 4], [4, 3], [5, 0]]

この場合の出力は 3 になります。友達グループは以下の3つに分かれるためです。

• グループ1：{0, 1, 5}
• グループ2：{2}
• グループ3：{3, 4}

解き方のアプローチ

この問題はグラフ理論における「連結成分」を数える問題そのものです。深さ優先探索（DFS）を使えば、効率よく解決できます。手順は以下の通りです。

1. nodes := 友達リストの要素数
2. visited := 要素数と同じ長さで False で埋めた訪問管理用リスト
3. ans := 0（グループ数のカウンター）
4. 関数 dfs(vertex) を定義する
   • visited[vertex] を True にする
   • friends[vertex] 内の各隣接ノード nei について、まだ訪問していなければ dfs(nei) を再帰呼び出しする
5. メイン処理では、ノード 0 から nodes-1 まで順にループし
   • まだ訪問していないノード i があれば dfs(i) を呼び出し、ans を1増やす
6. 最後に ans を返す

Pythonでの実装例

class Solution:
    def solve(self, friends):
        nodes = len(friends)
        visited = [False for _ in range(nodes)]
        ans = 0

        def dfs(vertex):
            visited[vertex] = True
            for nei in friends[vertex]:
                if not visited[nei]:
                    dfs(nei)

        for i in range(nodes):
            if not visited[i]:
                dfs(i)
                ans += 1

        return ans

ob = Solution()
friends = [
    [0, 1, 5],
    [1, 0],
    [2],
    [3, 4],
    [4, 3],
    [5, 0]
]
print(ob.solve(friends))

実行結果

入力：

[[0, 1, 5],
[1, 0],
[2],
[3, 4],
[4, 3],
[5, 0]]

出力：

3

計算量について

各ノードと各辺を高々一度ずつ訪問するため、時間計算量は O(N + E)（N は人数、E は友人関係の総数）、空間計算量は訪問管理リストと再帰スタックにより O(N) となります。大規模なデータでも効率的に動作する実装です。