[Python] ABC016 C

問題

C – 友達の友達

回答

友達の友達を数える

集合を使い、友達の友達を探索して数えます。

import collections

N, M = map(int, input().split())

adjacent_dict = collections.defaultdict(list)
for _ in range(M):
    a, b = map(int, input().split())
    adjacent_dict[a].append(b)
    adjacent_dict[b].append(a)

for i in range(1, N+1):
    # 自身の友達
    friends_i = adjacent_dict[i]
    
    # 友達の友達
    friends_of_friends = set()
    for j in friends_i:
        friends_j = adjacent_dict[j]
        friends_of_friends |= set(friends_j)
    
    # 友達の友達に自分自身がいる場合は除く
    if i in friends_of_friends:
        friends_of_friends.remove(i)
    # 友達の友達から自身の友達を除く
    diff_friends = friends_of_friends.difference(set(friends_i))
    
    print(len(diff_friends))

頂点を2つ選び共通の友人がいるか調べる

集合を使い、2つの頂点ごとに、共通の友人がいるかを調べます。

import collections

N, M = map(int, input().split())

adjacent_dict = collections.defaultdict(list)
for _ in range(M):
    a, b = map(int, input().split())
    adjacent_dict[a].append(b)
    adjacent_dict[b].append(a)


for i in range(1, N+1):
    # iの友達
    friends_i = set(adjacent_dict[i])

    count = 0
    for j in range(1, N+1):
        # 自身を除く
        if i == j:
            continue
        # jの友達
        frineds_j = set(adjacent_dict[j])
        
        # 友達は除く
        if i in frineds_j:
            continue
        
        mutual_friends_ij = friends_i & frineds_j
        # 自身、または直接の友達を除き、共通の友達がいる
        if mutual_friends_ij:
            count += 1

    print(count)

ダイクストラ法

ダイクストラ法を用いて、全頂点間の距離を求め、各頂点から距離2の頂点の数えます。

scipy

scipy.sparse.csgraph.dijkstra を使ってみます。

from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import dijkstra

N, M = map(int, input().split())

graph = [[0] * N for _ in range(N)]
for _ in range(M):
    a, b = map(int, input().split())
    graph[a-1][b-1] = 1
    graph[b-1][a-1] = 1

graph = csr_matrix(graph)

dist_matrix = dijkstra(graph)
for row in dist_matrix:
    counts = 0
    for dist in row:
        if dist == 2:
            counts += 1
    print(counts)

自分で実装

ダイクストラ法を自分で実装してみます。

import collections
import heapq

N, M = map(int, input().split())

# 隣接リスト
adjacent_dict = collections.defaultdict(list)
for _ in range(M):
    a, b = map(int, input().split())
    adjacent_dict[a].append(b)
    adjacent_dict[b].append(a)

# 全てのノードに対してダイクストラ法を適用
for i in range(1, N+1):
    # 全ての距離を最大値に指定
    distance_dict = collections.defaultdict(lambda: float('inf'))
    # 自身への距離は0
    distance_dict[i] = 0
    
    # (距離, 頂点)の優先度付きキュー
    priority_queue = []
    heapq.heappush(priority_queue, (0, i))

    # 訪問済みフラグ
    visited = collections.defaultdict(bool)

    while priority_queue:
        # 距離が最小であるノード
        cur_dist, cur_node = heapq.heappop(priority_queue)
        # 訪問済みフラグ
        if visited[cur_node]:
            continue
        visited[cur_node] = True

        # 隣接ノードとの距離
        for next_node in adjacent_dict[cur_node]:
            if visited[next_node]:
                continue
            new_dist = cur_dist + 1
            # 新しい距離が短い場合は距離を更新
            if distance_dict[next_node] > new_dist:
                distance_dict[next_node] = new_dist

                heapq.heappush(priority_queue, (new_dist, next_node))

    cnt = 0
    for val in distance_dict.values():
        if val == 2:
            cnt += 1
    print(cnt)

ワーシャルフロイド法

ワーシャルフロイド法を用いて、全頂点間の距離を求め、各頂点から距離2の頂点の数えます。

scipy

scipy.sparse.csgraph.floyd_warshall を使ってみます。

from scipy.sparse import csr_matrix
from scipy.sparse.csgraph import floyd_warshall

N, M = map(int, input().split())

graph = [[0] * N for _ in range(N)]
for _ in range(M):
    a, b = map(int, input().split())
    graph[a-1][b-1] = 1
    graph[b-1][a-1] = 1

graph = csr_matrix(graph)

dist_matrix = floyd_warshall(graph)

for row in dist_matrix:
    cnt = 0
    for dist in row:
        if dist == 2:
            cnt += 1
    print(cnt)

自分で実装

ワーシャルフロイド法を自分で実装してみます。

N, M = map(int, input().split())

# 距離行列の生成
inf = float('inf')
dist_matrix = [[inf] * N for _ in range(N)]
for i in range(N):
    dist_matrix[i][i] = 0
for _ in range(M):
    a, b = map(int, input().split())
    dist_matrix[a-1][b-1] = 1
    dist_matrix[b-1][a-1] = 1

# ワーシャルフロイド
for k in range(N):
    for i in range(N):
        for j in range(N):
            dist_matrix[i][j] = min(dist_matrix[i][j],
                                    dist_matrix[i][k] + dist_matrix[k][j])

for row in dist_matrix:
    cnt = 0
    for dist in row:
        if dist == 2:
            cnt += 1
    print(cnt)