[파이썬 | BOJ | 5567] 결혼식

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

출력

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

 

풀이

알고리즘 분류 : BFS

 

상근이는 항상 1번이므로, 1번을 시작점으로 BFS를 탐색하고, 그 결과 상근이로부터 거리가 1(친구) or 2(친구의 친구) 인것만 찾으면 된다.

 

input의 개수가 작으므로, queue를 파이썬의 list를 사용했다. input이 늘어나면 deque를 import해서 쓰는게 속도면에서 좋다.

 

import sys
read = sys.stdin.readline

def BFS(startV):
    q = [startV]
    length[startV] = 0
    while len(q):
        nowV = q.pop(0)
        for nextV in friends[nowV]:
            if length[nextV] > length[nowV]:
                length[nextV] = length[nowV] + 1
                q.append(nextV)


N = int(read())
L = int(read())
friends = [[] for _ in range(N+1)]      #인접 리스트
length = [999 for _ in range(N+1)]      #방문 체크 겸 depth
for _ in range(L):
    a, b = map(int, read().split())
    friends[a].append(b)
    friends[b].append(a)

BFS(1)
ans = 0
for i in range(1, N+1):
    if length[i] in [1, 2]:
        ans += 1
print(ans)