백준 5567 결혼식 파이썬 풀이

난이도 : 실버2

풀이일 : 04134

 

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

5567번: 결혼식

---

링크로 이동하기 귀찮은 분들을 위한 캡쳐

---

풀이 과정

처음에는 친구관계로 엮여 있지 않고 동떨어져 있는 사람을 제외한 모든 사람을 초대한다고 이해했는데, 힌트를 보니 상근이의 친구와 친구의 친구에 해당하는 사람을 세는 문제였다.

 

- 2차원 배열로 친구 관계를 저장하고, queue를 이용한 BFS로 동기들과 상근이의 관계를 리스트에 기록

- 상근이와 동기들의 관계는 상근이 본인 1, 직접친구 2, 이후로는 한 다리를 더 거쳐야 할 때마다 1씩 추가

- 최종적으로 리스트를 순회하며 숫자가 3 이하인 경우를 세서 출력

---

풀이 코드

import sys

# 초대하는 사람 판별 함수
def invite(num):
    people = 0
    visited = [0] * (n+1)
    visited[num] = 1
    queue = [num]
    while queue:
        i = queue.pop(0)
        for j in range(1, n+1):
            if friend[i][j] == 1 and visited[j] == 0:
                queue.append(j)
                visited[j] = visited[i] + 1

    # 친구의 친구까지 -> visited 숫자가 3이하
    for k in range(2, n+1):
        if visited[k] and visited[k] <= 3:
            people += 1

    return print(people)


n = int(sys.stdin.readline().strip())
m = int(sys.stdin.readline().strip())

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

# 2차원 배열에 친구 관계 저장
for _ in range(m):
    p, f = list(map(int, sys.stdin.readline().split()))
    friend[p][f] = friend[f][p] = 1

invite(1)

느낀점

동기 명단이 엄청나게 많아지거나 하면, 속도 제한에 걸릴 수 있을 것 같다.

향후 다시 풀어볼 때에는 동기와 상근이의 관계를 표현하는 배열 visited 숫자가 3 이상이면 중단하도록 수정해봐야겠다.