[ BOJ ] 5567 : 결혼식 ( SILVER 2 ) / Python

>> 문제 바로가기 (https://www.acmicpc.net/problem/5567)

문제

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

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

입력

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

출력

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

풀이 과정

문제를 잘 읽어야 한다. 필자는 자신의 친구의 친구의 친구를 초대하는 줄 알고 잘못 구현했었다..

자신의 친구, 그리고 자신의 친구의 친구를 초대해야 하므로 BFS 를 통해 친구 관계를 탐색하여 길이가 2 이하인 친구 노드 수를 세야 한다.

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

n = int(input().rstrip())
m = int(input().rstrip())

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

queue = deque()
queue.append(1)
visited = [0 for _ in range(n+1)]
visited[1] = 1

time = 0
while queue:
    size = len(queue)
    
    for s in range(size):
        now = queue.popleft()
        for next in list(graph[now].keys()):
            if visited[next] == 0:
                visited[next] = 1
                queue.append(next)
            
    if time == 1: break
    time += 1
print(sum(visited) - 1)