Dijkstra

ヒープに (距離, 頂点) のタプルではなく 距離 * N + 頂点 という整数を入れることで、PyPy での速度を上げています。そのため辺の重みは非負整数に限ります。

• dijkstra(g, s): 隣接リスト g[v] = [(to, cost), ...] と始点 $s$ から、dist（各頂点への最短距離）と prev（最短路木での親、始点と到達不能な頂点は $-1$）を返す。到達できない頂点の距離は 1 << 62
• restore_path(prev, t): 始点から $t$ までの最短路の頂点列を返す（ $t$ に到達できるときのみ使う）
• 計算量: $O((N + M) \log N)$

実数の重みを扱うときは、ヒープに (d, v) のタプルを入れる形に書き換えます。

コード

from heapq import heappush, heappop


def dijkstra(g, s):
    # g[v] = [(to, cost), ...]、cost は非負整数
    n = len(g)
    INF = 1 << 62
    dist = [INF] * n
    prev = [-1] * n
    dist[s] = 0
    hq = [s]
    while hq:
        dv = heappop(hq)
        d, v = divmod(dv, n)
        if d > dist[v]:
            continue
        for to, c in g[v]:
            nd = d + c
            if nd < dist[to]:
                dist[to] = nd
                prev[to] = v
                heappush(hq, nd * n + to)
    return dist, prev


def restore_path(prev, t):
    path = []
    while t != -1:
        path.append(t)
        t = prev[t]
    return path[::-1]