[코딩테스트] 백준 결혼식 파이썬(Python)

https://www.acmicpc.net/problem/5567

5567번: 결혼식

from collections import deque
n = int(input())
m = int(input())
graph = {i : [] for i in range(1, n+1)}

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

def dfs(v):
  visited = {}
  friends = set()
  q = deque()
  q.append((v, 0))
  while q:
    i, depth = q.popleft()
    if depth <= 2:
      friends.add(i)
    nodes = graph[i]
    visited[i]= True
    for node in nodes:
      if visited.get(node, False) == False:
        visited[node] = True
        q.append((node, depth+1))
  return friends - {1}


friends=dfs(1)
print(len(friends))

dfs를 이용하여 그래프의 깊이 depth를 확인하며 문제를 해결하였습니다.

 

문제에서 요구한 상근이 자신의 친구와 친구의 친구까지 초대할 수 있기 때문에 가능한 depth는 2입니다.
2이상의 관계는 친구의 친구의 친구이기 때문에 문제의 요구사항에 맞지 않습니다.

 

dfs 코드를 확인해보면 q에 노드 v와 깊이 depth를 각각 1, 0 을 초기에 저장합니다.

• v가 1인 이유는 상근이 자신이 1이기 때문입니다.

각 노드와 인접한 노드들을 확인하며 이미 방문처리된 노드는 확인하지 않고 방문처리 되어 있지 않은 노드만 확인합니다.
다음 노드를 확인하기 위하여 q에 저장할 때 depth + 1하여 그래프의 깊이를 하나 증가시켜 저장합니다.

 

popleft()로 노드를 추출할 때 깊이 depth도 같이 추출하게 되는데 이때 depth가 2보다 작거나 같다면
결혼식에 올 수 있는 조건에 부합하기 때문에 friends에 저장합니다.

 

 그래프 모두 확인했다면 friends에서 자기자신 1을 제거하여 반환합니다.

 

마지막으로 friends의 길이를 반환하여 문제를 해결합니다.