[백준] 5567. 구현_결혼식 - Python

문제

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

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

 

★ 시작하기 전에

친구의 친구만 구하는 경우라면 그저 조건문 2번을 활용하여 풀 수 있을 거라고 생각한다.

하지만 조건이 친구의 친구가 아니라 친구의 친구의 친구의 친구까지라면?

딱! dfs_bfs가 생각났다면 좋앗겠지만, 고민하고 고민해서 생각났다.

몇단계에 걸쳐서까지 친구인지를 구하고 싶어 q 내에 친구 관계 다리?를 넣어서 구현해보았다.

import sys
from collections import deque

input = sys.stdin.readline

n = int(input())
rel = int(input())
board = [[] for _ in range(n+1)]
for i in range(rel):
    a, b = map(int, input().split())
    board[a].append(b)
    board[b].append(a)
visited = [0]*(n+1)
#########################이상 visited 및 board를 입력 받음########################

q = deque()
# 초기값 1은 고정되어 있으며, 바뀌어도 값만 바꾸어주면 된다.
# 숫자 2는 '친구의 친구'이기 때문에 2다리로 시작하여 1씩 줄여갈 것이다.
q.append((1, 2))
visited[1] = 1
result = 0
# bfs를 활용해보자. q가 다 사라질 때까지 돌린다.
while q:
    now = q.popleft()
    point, cnt = now[0], now[1]
    # 현재 point사람의 친구관계를 불러온다.
    for i in board[point]:
        if visited[i] == 1:
            continue
        else:
        	# 가장 중요한 부분**이다. cnt가 0인 친구들까지는 q에 넣되, 그 밑으로 되면 넣지도 말자.
            if cnt-1 >= 0:
                q.append((i, cnt-1))
                visited[i] = 1
                result += 1
print(result)

※ 주의할 것

while q 내에서 bfs의 경우 같은 레벨의 항목들이 전부 들어가게 된다. 그것들까지 q를 돌아버리게 되기 때문에 조건에 맞지 않는 항목은 q에 넣지도 않고 while문이 종료되도록 구현하는 것이 중요해 보인다.