[백준] 5567번 - 결혼식 (파이썬)

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도 친구관계이다. 

출력

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

문제 탐색하기

 

• 상근이(1번)의 결혼식에 초대할 사람은 상근이의 직접적인 친구와 친구의 친구이다.
• 친구 관계는 양방향이며, 친구의 친구의 친구는 초대하지 않는다.
• 상근이의 친구 네트워크에서 최대 2단계 이내에 있는 사람들을 찾아야 한다.
• 초대할 사람의 수를 구하는 문제이다.

 

 

 

코드 설계하기

 

• 친구 관계를 저장하기 위해 인접 리스트를 사용한다.
• BFS를 이용하여 상근이로부터의 거리를 계산한다.
• 방문 여부와 거리를 기록하여 이미 방문한 노드를 재방문하지 않도록 한다.
• 상근이로부터 거리가 1 또는 2인 노드를 찾아 카운트한다.

 

 

시도 회차 수정사항

- 시간 초과가 발생

=> 입력된 값이 한 자리 수일 때를 처리할 수 있는 걸 모든 상황에서 제대로 작동하도록 조정

코드

from collections import deque

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)

visited = [0] * (n + 1)  # 방문 여부를 체크하기 위한 리스트
distance = [0] * (n + 1)  # 상근이로부터의 거리를 저장하는 리스트

queue = deque()
queue.append(1)  # 상근이부터 시작
visited[1] = 1  # 상근이 방문 체크

while queue:
    current = queue.popleft()
    for friend in graph[current]:
        if not visited[friend]:
            visited[friend] = 1  # 방문 체크
            distance[friend] = distance[current] + 1  # 거리 계산
            if distance[friend] <= 2:  # 거리 2 이하인 경우만 큐에 추가
                queue.append(friend)

# 상근이를 제외하고 거리 1 또는 2인 사람의 수를 계산
invite_count = sum(1 for i in range(2, n + 1) if 0 < distance[i] <= 2)

print(invite_count)  # 초대할 사람의 수 출력