[Python/파이썬] 백준 5567번 결혼식

5567번: 결혼식

 

문제

상근이는 자신의 결혼식에 학교 동기 중 자신의 친구와 친구의 친구를 초대하기로 했다. 상근이의 동기는 모두 N명이고, 이 학생들의 학번은 모두 1부터 N까지이다. 상근이의 학번은 1이다.

상근이는 동기들의 친구 관계를 모두 조사한 리스트를 가지고 있다. 이 리스트를 바탕으로 결혼식에 초대할 사람의 수를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 상근이의 동기의 수 n (2 ≤ n ≤ 500)이 주어진다. 둘째 줄에는 리스트의 길이 m (1 ≤ m ≤ 10000)이 주어진다. 다음 줄부터 m개 줄에는 친구 관계 ai $b_i$가 주어진다. (1 ≤ ai < $b_i$ ≤ n) ai와 $b_i$가 친구라는 뜻이며, $b_i$와 $a_i$도 친구관계이다. 

출력

첫째 줄에 상근이의 결혼식에 초대하는 동기의 수를 출력한다.

코드

from collections import deque
import sys
input = sys.stdin.readline

n = int(input())
m = int(input())
graph = [[] for _ in range(n+1)]
for _ in range(m):
    a, b = map(int, input().split())
    graph[a].append(b)
    graph[b].append(a)

q = deque([(1, 0)])
visited = [False] * (n+1)
visited[1] = True
ans = 0
while q:
    now, cnt = q.popleft()
    if cnt == 2:    # 친구의 친구까지만 초대
        continue
    for nx in graph[now]:
        if not visited[nx]:
            q.append((nx, cnt+1))
            visited[nx] = True
            ans += 1
print(ans)

상근이의 친구관계를 그래프로 나타냈을때, 상근이(1)로부터 거리가 2 이하로 연결되어 있는 노드의 개수를 세면 된다. 이를 위해 BFS를 이용하였다.