[백준] 5567번 결혼식 Python

📃문제

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

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

📃입력

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

📃출력

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

📃예제 입력 1

6
5
1 2
1 3
3 4
2 3
4 5

📃예제 출력 1

3

📃예제 입력 2

6
5
2 3
3 4
4 5
5 6
2 5

📃예제 출력 2

0

📃힌트

예제 1의 경우 2와 3은 상근이의 친구이다. 또, 3과 4는 친구이기 때문에, 4는 상근이의 친구의 친구이다. 5와 6은 친구도 아니고, 친구의 친구도 아니다. 따라서 2, 3, 4 3명의 친구를 결혼식에 초대한다.

✏️문제 탐색하기

상근이가 자신의 결혼식에 초대할 사람의 수를 구하는 문제.

주어진 조건:

• 상근이의 학번은 1번.
• n: 동기의 수 ( 2 <= n <= 500 )
• m: 친구 관계의 수 ( 1 <= m <= 10000 )

✏️알고리즘 선택

BFS알고리즘으로 선택
모든 노드와 간선을 한 번씩 방문하므로 O(N * M)의 시간복잡도를 가짐.

✏️코드 설계하기

1. 동기 수 n, 친구 관계 수 m입력.
2. 친구 관계를 인접 리스트로 표현.
3. 시작 노드 1에서 BFS 실행.
4. 깊이 2까지만 탐색하여 상근이의 친구와 친구의 친구를 찾음.
5. visited 배열로 중복 방문 방지.
6. BFS 결과로 방문한 노드 수에서 상근이 자신을 제외한 값 출력.

✏️코드 구현

import sys

from collections import deque


N = int(sys.stdin.readline())
M = int(sys.stdin.readline())

list_ = [[] for _ in range(N+1)]
for _ in range(M):
    a, b = list(map(int, sys.stdin.readline().split()))
    list_[a].append(b)
    list_[b].append(a)


def bfs(start):
    visited = [False] * (N + 1)
    queue = deque([(start, 0)])
    visited[start] = True
    count = 0

    while queue:
        node, depth = queue.popleft()
        if depth > 2:
            continue
        if depth > 0:
            count += 1

        for v in list_[node]:
            if not visited[v]:
                visited[v] = True
                queue.append((v, depth+1))
    return count


print(bfs(1))