[ BOJ / Python ] 1389번 케빈 베이컨의 6단계 법칙

황승환·2022년 3월 2일
0

Python

목록 보기
215/498


이번 문제는 바로 이전에 풀었던 문제와 거의 풀이가 유사하다. 다익스트라 알고리즘을 함수로 구현하고, 모든 노드에서부터의 다익스트라 함수를 모두 호출한 후, 반환되는 거리 리스트의 원소들의 총합을 한 리스트에 모은 후 가장 작은 값의 인덱스를 출력하도록 접근하였다.

  • n, m을 입력받는다.
  • 그래프로 사용할 리스트 graph를 2차원 리스트로 선언한다.
  • m번 반복하는 for문을 돌린다.
    -> a, b를 입력받는다.
    -> graph[a]에 b를 넣는다.
    -> graph[b]에 a를 넣는다.
  • search함수를 start를 인자로 갖도록 선언한다.
    -> 최댓값을 INF에 저장한다.
    -> dist를 INF n+1개 채운 리스트로 선언한다.
    -> q를 최소 힙으로 선언한다.
    -> q에 (0, start)를 넣는다.
    -> q가 존재하는 동안 반복하는 while문을 돌린다.
    --> q에서 cost, cur을 추출한다.
    --> 만약 cost가 dist[cur]보다 클 경우, 다음 반복으로 넘어간다.
    --> graph[cur]을 순회하는 nxt에 대한 for문을 돌린다.
    ---> 만약 dist[nxt]가 cost+1보다 클 경우,
    ----> dist[nxt]를 cost+1로 갱신한다.
    ----> q에 (cost+1, nxt)를 넣는다.
    -> dist를 반환한다.
  • answers를 0 n+1개로 채운 리스트로 선언한다.
  • 1부터 n까지 반복하는 i에 대한 for문을 돌린다.
    -> 임시 변수 tmp에 search(i)의 반환값을 저장한다.
    -> answers[i]tmp[1:]의 총합을 더한다.
  • answers에서 가장 작은 값의 인덱스를 출력한다.

Code

import heapq
import sys
n, m=map(int, input().split())
graph=[[] for _ in range(n+1)]
for _ in range(m):
    a, b=map(int, input().split())
    graph[a].append(b)
    graph[b].append(a)
def search(start):
    INF=sys.maxsize
    dist=[INF for _ in range(n+1)]
    q=[]
    heapq.heappush(q, (0, start))
    while q:
        cost, cur=heapq.heappop(q)
        if cost>dist[cur]:
            continue
        for nxt in graph[cur]:
            if dist[nxt]>cost+1:
                dist[nxt]=cost+1
                heapq.heappush(q, (cost+1, nxt))
    return dist
answers=[0]*(n+1)
for i in range(1, n+1):
    tmp=search(i)
    answers[i]+=sum(tmp[1:])
print(answers.index(min(answers[1:])))

profile
꾸준함을 꿈꾸는 SW 전공 학부생의 개발 일기

0개의 댓글