[BOJ/백준] 5567번 결혼식 (Python 파이썬)

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

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 

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명의 친구를 결혼식에 초대한다.

 

---

풀이

우선, 주어진 input으로 친구 관계를 나타내는 그래프를 완성시켜야 합니다. input으로는 친구 관계가 한 방향으로만 주어지지만, 실제로는 그래프내에서 양방향으로 존재하므로 아래와 같이 구현합니다.

graph = [[0] * (n+1) for i in range(n+1)] # 1번 인덱스부터 사용

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

 

이제, BFS을 통해서 서로 아는 관계(친구 관계로 연결된 관계)을 구하면 되는데, 여기서 문제는 "상근이의 직접 친구"와 "상근이의 친구의 친구"까지만 허용되는 점을 고려해야 합니다.

이를 탐색의 관점에서 바라본다면, BFS의 탐색 depth를 3까지만 허용한다는 의미입니다. 

그럼, 앞선 문제를 풀었듯이 방문 여부를 관리할 때 누적합을 통해서 탐색의 depth을 넣어주고, 최종적으로 탐색이 종료되었을 때 depth가 0보다 크고 3이하인 친구들의 수를 세면 됩니다!

이를 구현한 코드는 아래와 같습니다.

 

Code

import sys
from collections import deque
input = sys.stdin.readline
n = int(input())
m = int(input())
graph = [[0] * (n+1) for i in range(n+1)] # 1번 인덱스부터 사용
visited = [0 for _ in range(n+1)]

for i in range(m):
    a,b = map(int,input().split())
    graph[a].append(b)
    graph[b].append(a)
    
# 1번 인덱스 -> 상근이 -> 상근이의 친구들
# 상근이 친구들의 친구들을 탐색한다. -> 이때 이미 초대할 사람인 지를 확인하여 중복되지 않도록 한다.
# 탐색 깊이는 최대 2이다. -> visited의 값을 누적합으로 정해서 탐색의 depth를 표현
# -> 어떤 사람의 visited의 값은 최대 3이다. (1: 본인, 2: 직접 친구, 3: 친구의 친구)

def BFS(x):
  queue = deque([x])
  visited[x] = 1
  
  while queue:
    friend_num = queue.popleft()
    for people_num in graph[friend_num]:
      if visited[people_num] == 0:
        queue.append(people_num)
        visited[people_num] = visited[friend_num] + 1


def solve():
  BFS(1)
  result = 0
  for i in range(2, n+1):
    if 0 < visited[i] <=3:
       result += 1
  print(result)

solve()