Python – 표준 모듈(heapq편)

• heappush

heapq.heappush(q, 1)
heapq.heappush(q, 2)
heapq.heappush(q, 3)

print(q)
# [1, 2, 3]

• heappop

heapq.heappop(q)
# 1

• heappushpop

heapq.heappushpop(q, 4)
# 2
print(q)
# [3, 4]

• heapify – 기본적으로 최소 힙을 제공 (참고 – 힙 알고리즘)

items = [4,1,2,3,5,6]

q = []
for item in items:
    heapq.heappush(q, item)
    print(q)

[4]
[1, 4]
[1, 4, 2]
[1, 3, 2, 4]
[1, 3, 2, 4, 5]
[1, 3, 2, 4, 5, 6]

q = [4,1,2,3,5,6]
heapq.heapify(q)
print(q)
# [1, 3, 2, 4, 5, 6]

• heapq 활용해보기 – 백준 1504번 (참고 – 다익스트라 알고리즘)

문제 (https://www.acmicpc.net/problem/1504)

방향성이 없는 그래프가 주어진다. 세준이는 1번 정점에서 N번 정점으로 최단 거리로 이동하려고 한다. 또한 세준이는 두 가지 조건을 만족하면서 이동하는 특정한 최단 경로를 구하고 싶은데, 그것은 바로 임의로 주어진 두 정점은 반드시 통과해야 한다는 것이다.

세준이는 한번 이동했던 정점은 물론, 한번 이동했던 간선도 다시 이동할 수 있다. 하지만 반드시 최단 경로로 이동해야 한다는 사실에 주의하라. 1번 정점에서 N번 정점으로 이동할 때, 주어진 두 정점을 반드시 거치면서 최단 경로로 이동하는 프로그램을 작성하시오.

import sys
import heapq
N, E = map(int, sys.stdin.readline().split())
INF = int(1e9)

graph = {}
for i in range(E):
    # 입력값 받기
    p1, p2, v = map(int, sys.stdin.readline().split())
    # 그래프 그리기
    if graph.get(p1):
        graph.get(p1)[p2] = v
    else:
        graph[p1] = {p2 : v}

    if graph.get(p2):
        graph.get(p2)[p1] = v
    else:
        graph[p2] = {p1 : v}
v1, v2 = map(int, sys.stdin.readline().split())
q = []

def dijkstra(start, end):
    # 비교값 저장
    distance = [INF] * (N+1)
    distance[start] = 0

    q = [(0, start)]

    while q:
        dist, now = heapq.heappop(q)
        for i in graph[now]:
            cost = dist + graph[now][i] # 현재까지의 이동 비용에서 다음 경로로 이동하는 비용 추가
            if cost < distance[i]: # 이전에 저장된 경로까지의 비용과 현재 이동 경로의 비용 중 최소 비용 비교
                distance[i] = cost
                heapq.heappush(q, (cost, i))

    return distance[end]

path1 = dijkstra(1, v1) + dijkstra(v1, v2) + dijkstra(v2, N)
path2 = dijkstra(1, v2) + dijkstra(v2, v1) + dijkstra(v1, N)
result = min(path1, path2)

if result >= INF:
    print(-1)
else:
    print(result)