Shortest Path in DAGDAG 最短路徑

from collections import deque

INF = float("inf")


def topo_order(n, adj):
    """Kahn 拓撲排序；輸出少於 n 個節點就代表圖裡有環"""
    indeg = [0] * n
    for u in range(n):
        for v, _ in adj[u]:
            indeg[v] += 1
    q = deque(i for i in range(n) if indeg[i] == 0)
    order = []
    while q:
        u = q.popleft()
        order.append(u)
        for v, _ in adj[u]:
            indeg[v] -= 1
            if indeg[v] == 0:
                q.append(v)
    if len(order) != n:
        raise ValueError("圖裡有環，不是 DAG")
    return order


def dag_shortest(n, edges, src):
    """DAG 單源最短路徑，允許負權。回傳 (dist, parent)，O(V + E)"""
    adj = [[] for _ in range(n)]
    for u, v, w in edges:
        adj[u].append((v, w))
    dist, parent = [INF] * n, [-1] * n
    dist[src] = 0
    for u in topo_order(n, adj):
        if dist[u] == INF:                      # 走不到的點不能拿來鬆弛
            continue
        for v, w in adj[u]:
            if dist[u] + w < dist[v]:
                dist[v], parent[v] = dist[u] + w, u
    return dist, parent


def critical_path(durations, deps):
    """專案排程：durations[i] 是工作 i 的天數，deps 是 (前置, 後續) 的清單。
    回傳 (總工期, 每項工作的浮動時間)。浮動時間 0 的工作就在關鍵路徑上"""
    n = len(durations)
    adj = [[] for _ in range(n)]
    for u, v in deps:
        adj[u].append((v, durations[u]))
    order = topo_order(n, adj)
    earliest = [0] * n                          # 最早開始 = 從起點出發的最長路徑
    for u in order:
        for v, w in adj[u]:
            earliest[v] = max(earliest[v], earliest[u] + w)
    total = max(earliest[i] + durations[i] for i in range(n))
    latest = [total - durations[i] for i in range(n)]   # 最晚開始：再晚就會拖到總工期
    for u in reversed(order):                   # 反向拓撲順序，由後往前推
        for v, w in adj[u]:
            latest[u] = min(latest[u], latest[v] - w)
    return total, [latest[i] - earliest[i] for i in range(n)]


if __name__ == "__main__":
    R, S, T, X, Y, Z = range(6)
    edges = [(R, S, 5), (R, T, 3), (S, T, 2), (S, X, 6), (T, X, 7), (T, Y, 4), (T, Z, 2), (X, Y, -1), (X, Z, 1), (Y, Z, -2)]
    dist, parent = dag_shortest(6, edges, S)
    print(dist)                                 # [inf, 0, 2, 6, 5, 3]
    path, v = [], Z
    while v != -1:
        path.append("RSTXYZ"[v])
        v = parent[v]
    print("".join(reversed(path)))              # SXYZ：6 + (-1) + (-2) = 3

    # 開工、規格、後端、前端、文件、測試、上線
    dur = [0, 3, 6, 4, 2, 3, 0]
    deps = [(0, 1), (1, 2), (1, 3), (1, 4), (2, 5), (3, 5), (4, 6), (5, 6)]
    print(critical_path(dur, deps))             # (12, [0, 0, 0, 2, 7, 0, 0])