[BOJ] 5567번 결혼식 파이썬

문제

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

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

입력

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

출력

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

풀이

#graph

처음에 플로이드워셜이 생각나서 그 방식으로 풀었는데(단, 친구의 친구만 구하는 거기 때문에 상근이의 관계에 대해서만 체크. 전체 조사하면 시간초과됨) 그냥 반복문 돌면서 방문 체크해주면서 구현하는 게 더 효율적일 거 같다.

코드

n = int(input())
m = int(input())
INF = 1e9
count = 0

arr = [[INF]*(n+1) for _ in range(n+1)]
for i in range(1,n+1):
    arr[i][i]=0

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

count = 0
for i in range(n+1):
    for b in range(n+1):
        arr[1][b]=min(arr[1][i]+arr[i][b], arr[1][b])

for i in range(n+1):
    if 1<=arr[1][i]<=2:
        count+=1

print(count)

코드2

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

arr = [[] for _ in range(n+1)]
visited = [0] * (n+1)

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

count = 0
visited[1] = 1
for i in arr[1]:
    if not visited[i]:
        visited[i] = 1
        count+=1
    for j in arr[i]:
        if not visited[j]:
            visited[j] = 1
            count+=1

print(count)