[백준 5567번] 결혼식 - python

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

📌 문제 탐색하기

목표

• 친구의 친구까지 결혼식에 초대할 때 결혼식에 초대하는 동기의 수를 출력하자.

입력값

• 첫째 줄에 상근이의 동기수 n을 입력한다. (2 ≤ n ≤ 500)
• 둘째 줄에 리스트의 길이 m을 입력한다. (1 ≤ m ≤ 10000)
• m개의 줄에는 친구 관계 a, b가 주어진다.  (1 ≤ a < b ≤ n)
• ex)

6
5
1 2
1 3
3 4
2 3
4 5

출력값

• 상근이의 결혼식에 초대하는 동기의 수를 출력한다.
• ex)
• 3

📌 코드 설계하기

해결 아이디어

• 이번 문제는 그래프 탐색을 하는 문제이다. dfs에서 했듯 모든 관계를 배열에 삽입하여 친구 관계를 저장한다.
• 그리고 dfs를 2번만 돌아서 친구의 친구까지만 count할 수 있도록 하자.
• 이렇게 코드를 짜고 있었는데 훨씬 쉬운 방법이 생각났다.
• 상근이의 수는 1로 정해져 있고 그러면 상근이의 직속 친구들은 friend[1]에 모두 저장되어있다. 따라서 friend[1]의 각 요소에 접근해서 그들의 친구들의 수를 세면 된다.
• 물론 여기서 가장 중요한게 중복을 관리하는 것이다. 일단 직속 친구를 invite배열에 삽입하고 친구의 친구들을 중복없이 삽입하면 문제를 쉽게 풀 수 있다!!

📌 정답 코드

import sys

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

friend = [[] for _ in range (n+1)]
for i in range (m):
    a, b = map(int, sys.stdin.readline().split())
    friend[a].append(b)
    friend[b].append(a)

friends = 0
invite = []
for fri in friend[1]:
    if fri not in invite:
        invite.append(fri)
    for i in friend[fri]:
        if i == 1:
            continue
        if i not in invite:
            invite.append(i)

print(len(invite))

시간 복잡도 구해보기

• for문
  • friend[1]의 길이가 최대 n-1 → O(N)
  • friend[fri]의 길이가 최대 n-1 → O(N)
  • 중복 체크 → O(N)
  • 모두 곱하면 O($N^3$)

최종 시간 복잡도: O($N^3$)

너무 크다..! 줄여보자

📌 다른 풀이

import sys

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

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

for i in range (m):
    a, b = map(int, sys.stdin.readline().split())
    friend[a].append(b)
    friend[b].append(a)

invite = set()

for fri in friend[1]:
    invite.add(fri)
    for i in friend[fri]:
        if i != 1:
            invite.add(i)

print(len(invite))

• python의 set은 집합 자료형이다.
• 중복을 허용하지 않는다.
• 원소 존재 여부 확인이 빠르다. O(1)으로 가능하다.
• 원소들이 순서없이 저장되기 때문에 list처럼 인덱스로 접근할 수 없다.
• 집합이라서 중복이 자동으로 제거되기 때문에 시간 복잡도가 $O(N^2)$으로 줄어든다.