lessons/互动程序/06-最小与最大生成树/00-只修够用的路.py

#!/usr/bin/env python3
"""课前零食：追踪 Kruskal 和并查集的每一次合并。"""

from __future__ import annotations


DEFAULT = [
    ("A", "B", 1), ("D", "E", 1), ("A", "C", 2), ("C", "D", 2),
    ("B", "C", 3), ("B", "D", 4), ("C", "E", 5),
]


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


def parse_edges(raw: str) -> list[tuple[str, str, int]]:
    result: list[tuple[str, str, int]] = []
    for token in raw.replace("，", ",").split(","):
        parts = [part.strip() for part in token.split("-")]
        if len(parts) != 3 or not parts[0] or not parts[1]:
            raise ValueError("每条边应形如 A-B-3")
        result.append((parts[0], parts[1], int(parts[2])))
    return result


class DSU:
    def __init__(self, vertices: set[str]) -> None:
        self.parent = {vertex: vertex for vertex in vertices}
        self.size = {vertex: 1 for vertex in vertices}

    def find(self, vertex: str) -> str:
        if self.parent[vertex] != vertex:
            self.parent[vertex] = self.find(self.parent[vertex])
        return self.parent[vertex]

    def union(self, left: str, right: str) -> bool:
        left, right = self.find(left), self.find(right)
        if left == right:
            return False
        if self.size[left] < self.size[right]:
            left, right = right, left
        self.parent[right] = left
        self.size[left] += self.size[right]
        return True

    def groups(self) -> list[list[str]]:
        result: dict[str, list[str]] = {}
        for vertex in self.parent:
            result.setdefault(self.find(vertex), []).append(vertex)
        return [sorted(group) for group in result.values()]


def main() -> None:
    raw = ask("边表 A-B-3,B-C-1（回车用预设）： ")
    try:
        edges = DEFAULT if not raw else parse_edges(raw)
    except (ValueError, TypeError) as error:
        print(f"输入格式有问题：{error}；使用预设。")
        edges = DEFAULT
    vertices = {vertex for edge in edges for vertex in edge[:2]}
    dsu = DSU(vertices)
    chosen: list[tuple[str, str, int]] = []
    total = 0
    print("\n--- 按权排序后 ---")
    for u, v, weight in sorted(edges, key=lambda edge: edge[2]):
        before = (dsu.find(u), dsu.find(v))
        if dsu.union(u, v):
            chosen.append((u, v, weight))
            total += weight
            print(f"选 {u}-{v} ({weight})：分量 {before[0]} 与 {before[1]} 合并 -> {dsu.groups()}")
        else:
            print(f"跳过 {u}-{v} ({weight})：两端都在分量 {before[0]}，会闭合环")
    print("\n选中的边：", chosen)
    if len(chosen) == len(vertices) - 1:
        print(f"MST 总权={total}，恰好选了 {len(vertices) - 1} 条边。")
    else:
        print("图不连通，只得到最小生成森林：", dsu.groups())


if __name__ == "__main__":
    main()