[백준] 1238 파티 파이썬

오늘의 학습 키워드

다익스트라

문제

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

코드

정답 코드

import heapq
INF = int(1e9)
n, m, x = map(int, input().split())
graph = [[] for i in range(n+1)]
for _ in range(m):
    a, b, l = map(int, input().split())
    graph[a].append((b, l))

def dijkstra(start, end):
    q = []
    distance = [INF] * (n+1)
    heapq.heappush(q, (0, start))
    distance[start] = 0
    while q:
        dist, node = heapq.heappop(q)
        if distance[node] < dist:
            continue
        for new in graph[node]:
            new_dist = distance[node] + new[1]
            if new_dist < distance[new[0]]:
                distance[new[0]] = new_dist
                heapq.heappush(q, (new_dist, new[0]))
    return distance[end]

answer = 0

for i in range(1, n+1):
    if i == x:
        continue
    answer = max(answer, dijkstra(i, x) + dijkstra(x, i))

print(answer)

 

해당 문제는 n명의 학생이 x 마을에 단방향 도로를 사용해서 왕복하는 최단 거리를 계산하고, 학생들 중 가장 많은 거리를 간 학생의 이동 값을 구하는 문제이다. 출발지와 목적지가 정해져 있고 최단 거리를 계산한다는 점에서 '다익스트라' 알고리즘이다.

학생의 이동거리는 집 => x마을 => 집 임으로 다익스트라 계산을 통해 모든 노드 => x, x => 모든 노드 값을 구해야 한다.

먼저 다익스트라 코드를 구현한다. 아래 코드는 기본 다익스트라 코드이다. 

 

기본 다익스트라 코드(백준 1753 최단경로)

import heapq
import sys
input = sys.stdin.readline
INF = int(1e9)

v, e = map(int, input().split())
k = int(input())
graph = [[] for i in range(v+1)]
for _ in range(e):
    a, b, c = map(int, input().split())
    graph[a].append((b, c))

distance = [INF] * (v+1)

def dijkstra(start):
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0
    while q:
        dist, b = heapq.heappop(q)
        if distance[b] < dist:
            continue
        for i in graph[b]:
            cost = dist + i[1]
            if cost < distance[i[0]]:
                distance[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))

dijkstra(k)

for i in range(1, v+1):
    if distance[i] == INF:
        print("INF")
    else:
        print(distance[i])

 

해당 코드를 응용해서 다익스트라 알고리즘을 짠 다음, 마을의 크기만큼 for문을 돌면서 각 마을에서 x마을까지의 값과, x마을에서 각 마을의 값을 각자 구해 이를 더해준다. 이후 최대값을 지속적으로 업데이트 해주며 이동한 값이 가장 많은 학생의 이동값을 갱신해준다.