[Python] 백준 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도 친구관계
• 출력
  • 첫째 줄에 상근이의 결혼식에 초대하는 동기의 수를 출력
• 시간 제한: 1초
• 메모리 제한: 128 MB

 

"""
입출력 예시)

(예제 입력 1) 
n = 6
m = 5
1 2
1 3
3 4
2 3
4 5
    -> 3

(예제 입력 2)
n = 6
m = 5
2 3
3 4
4 5
5 6
2 5
    -> 0
"""

# (예제 입력 1)
# n = 6
# m = 5
# 1 2
# 1 3
# 3 4
# 2 3
# 4 5

# -> 2와 3은 상근이 친구
# -> 3과 4는 친구이기 때문에 4는 상근이의 친구의 친구
# -> 5, 6은 상근이 친구도 x, 친구의 친구도 x

# => 3명의 친구(2, 3, 4)를 결혼식에 초대

 

 

## 의사코드 ##

# 자신의 친구와 친구의 친구를 초대 -> BFS (형제 노드 탐색)

# 입력 값 받기
#   sys.stdin.readline()

# 1) BFS로 탐색하면서 거리를 구함
#   -> deque 
#   visited 큐 대신 distance 큐(친구에 대한 거리) 사용 

# 2) 친구의 친구 까지만 초대 -> 거리가 2보다 작거나 같은 경우인 노드의 수를 카운트
# 3) 카운트한 수 - 1(자기 자신)

 

 

 

첫 번째 시도

 

import sys
from collections import deque

n = int(input()) # 동기의 수
m = int(input()) # 리스트의 길이 m

# 그래프에 친구 관계 입력
graph = [[] for _ in range(n+1)]
for _ in range(m):
    a, b = map(int, sys.stdin.readline().split())
    graph[a].append(b)
    graph[b].append(a)

# 모든 친구에 대한 거리 초기화
distance = [-1] * (n+1)
# 시작점까지의 거리는 0으로 설정
distance[1] = 0

# BFS
q = deque([1]) # 상근이(1)부터 출발 
while q:
    x = q.popleft() # 현재 노드
    # 현재 노드에서 이동할 수 있는 모든 노드 확인
    for y in graph[x]:
        # 방문하지 않은 노드인 경우
        if distance[y] == -1:
            distance[y] = distance[x] + 1 # 방문처리 (현재 노드의 distance + 1)
            q.append(y)

result = 0 # 초대할 사람의 수 
for i in range(1, n+1): # 1부터 n까지 
    # i가 방문처리 되어 있고, distance[i]가 2 이하인 경우
    if distance[i] != -1 and distance[i] <= 2:
        result += 1 # 초대 

print(result - 1) # 자기 자신 제외

-> 런타임 에러

 

 

 

통과한 코드

 

import sys
from collections import deque
input = sys.stdin.readline

n = int(input()) # 동기의 수
m = int(input()) # 리스트의 길이 m

# 그래프에 친구 관계 입력
graph = [[] for _ in range(n+1)]
for _ in range(m):
    a, b = map(int, input().split())
    graph[a].append(b)
    graph[b].append(a)

# 모든 친구에 대한 거리 초기화
distance = [-1] * (n+1)
# 시작점까지의 거리는 0으로 설정
distance[1] = 0

# BFS
q = deque([1]) # 상근이(1)부터 출발 
while q:
    x = q.popleft() # 현재 노드
    # 현재 노드에서 이동할 수 있는 모든 노드 확인
    for y in graph[x]:
        # 방문하지 않은 노드인 경우
        if distance[y] == -1:
            distance[y] = distance[x] + 1 # 방문처리 (현재 노드의 distance + 1)
            q.append(y)

result = 0 # 초대할 사람의 수 
for i in range(1, n+1): # 1부터 n까지 
    # i가 방문처리 되어 있고, distance[i]가 2 이하인 경우
    if distance[i] != -1 and distance[i] <= 2:
        result += 1 # 초대 

print(result - 1) # 자기 자신 제외

-> 모든 입력을 sys.stdin.readline()으로 받는 것으로 바꿈

 

 

• graph 입력 예시 

n=6
m=5

graph = [[] for _ in range(n+1)]
graph 
"""
[[], [], [], [], [], [], []]
"""

graph[1].append(2)
graph[2].append(1)
graph
"""
[[], [2], [1], [], [], [], []]
"""