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

문제 출처 : 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도 친구관계이다. 

출력

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

 

내 풀이

n = int(input())
m = int(input())

arr = dict()
friend = []


for i in range(m):
    x, y = map(int, input().split())
    if x not in arr:
        arr[x] = list()
    if y not in arr:
        arr[y] = list()
    arr[x].append(y)
    arr[y].append(x)

    
for i in arr[1]:
    friend.append(i)
    if i in arr:
        for j in arr[i]:
            friend.append(j)
            
friend = set(friend)

print(len(friend) - 1)

처음에는 한 방향으로만 고려를 해주게 되면 될거라고 생각했던게 시간을 좀 많이 잡아먹었다. 

하지만 생각해보니 양방향 모두 고려를 해주어야했다. 

그래프를 만드는 것 보다 dict를 만들어 각 사람들의 key에 친구 목록을 item으로 저장해주고 

친구의 친구까지만 고려를 하기 때문에 1번에 저장된 사람들의 친구까지만 체크를 해주면서 friend 목록에 넣고

자기 자신과 중복된 경우를 제외하기 위해 set함수를 한 번 써준후 len(friend)에서 1을 뺀 값을 출력하도록 했다.