lessons/互动程序/05-最短路与类最短路/00-从乱走到松弛.py

#!/usr/bin/env python3
"""课前零食：比较 BFS、0-1 BFS 与 Dijkstra 的取点顺序。"""

from __future__ import annotations

from collections import deque
import heapq


EDGES = [
    ("S", "A", 1), ("S", "B", 0), ("B", "A", 0),
    ("A", "T", 1), ("B", "T", 1), ("S", "C", 1), ("C", "T", 1),
]


def ask(prompt: str, default: str) -> str:
    try:
        return input(prompt).strip() or default
    except EOFError:
        return default


def graph() -> dict[str, list[tuple[str, int]]]:
    result: dict[str, list[tuple[str, int]]] = {}
    for u, v, weight in EDGES:
        result.setdefault(u, []).append((v, weight))
        result.setdefault(v, [])
    return result


def zero_one_bfs(adj: dict[str, list[tuple[str, int]]]) -> None:
    dist = {u: float("inf") for u in adj}
    dist["S"] = 0
    work: deque[str] = deque(["S"])
    print("\n--- 0-1 BFS ---")
    while work:
        u = work.popleft()
        print(f"取出 {u}，dist={dist[u]}，双端队列={list(work)}")
        for v, weight in adj[u]:
            candidate = dist[u] + weight
            if candidate < dist[v]:
                old = dist[v]
                dist[v] = candidate
                if weight == 0:
                    work.appendleft(v)
                    place = "队首（0 边）"
                else:
                    work.append(v)
                    place = "队尾（1 边）"
                print(f"  松弛 {u}->{v}({weight})：{old} -> {candidate}，放{place}")
    print("最终距离：", dist)


def dijkstra(adj: dict[str, list[tuple[str, int]]]) -> None:
    dist = {u: float("inf") for u in adj}
    dist["S"] = 0
    heap: list[tuple[int, str]] = [(0, "S")]
    print("\n--- Dijkstra ---")
    while heap:
        current, u = heapq.heappop(heap)
        if current != dist[u]:
            print(f"弹出旧条目 ({current}, {u})，当前已是 {dist[u]}，跳过。")
            continue
        print(f"定案候选 {u}，距离 {current}")
        for v, weight in adj[u]:
            candidate = current + weight
            if candidate < dist[v]:
                old = dist[v]
                dist[v] = candidate
                heapq.heappush(heap, (candidate, v))
                print(f"  松弛 {u}->{v}({weight})：{old} -> {candidate}，推入堆")
    print("最终距离：", dist)


def bfs_unit_edges(adj: dict[str, list[tuple[str, int]]]) -> None:
    dist = {u: -1 for u in adj}
    dist["S"] = 0
    work: deque[str] = deque(["S"])
    print("\n--- 把每条边都当 1 的 BFS（仅作对照）---")
    while work:
        u = work.popleft()
        for v, _ in adj[u]:
            if dist[v] == -1:
                dist[v] = dist[u] + 1
                work.append(v)
    print("按边数得到：", dist)
    print("它把 S->A 当作 1 步最短，却看不见 S->B->A 的总权为 0。")


def main() -> None:
    print("预设有向边：", ", ".join(f"{u}->{v}({w})" for u, v, w in EDGES))
    choice = ask("选择 b=BFS 对照、z=0-1 BFS、d=Dijkstra（回车为 z）： ", "z").lower()
    adj = graph()
    if choice == "b":
        bfs_unit_edges(adj)
    elif choice == "d":
        dijkstra(adj)
    else:
        zero_one_bfs(adj)


if __name__ == "__main__":
    main()