lessons/互动程序/06-最小与最大生成树/00-把合并史写成树.py

#!/usr/bin/env python3
"""课前零食：把 Kruskal 的成功合并打印成一棵历史树。"""

from __future__ import annotations


EDGES = (
    (1, 2, 2),
    (3, 4, 3),
    (2, 3, 5),
    (4, 5, 6),
    (1, 5, 9),
)


class ReconstructionTree:
    def __init__(self) -> None:
        self.dsu_parent: dict[int, int] = {vertex: vertex for vertex in range(1, 6)}
        self.tree_parent: dict[int, int | None] = {vertex: None for vertex in range(1, 6)}
        self.children: dict[int, list[int]] = {vertex: [] for vertex in range(1, 6)}
        self.value: dict[int, int | None] = {vertex: None for vertex in range(1, 6)}
        self.logs: list[str] = []
        self.next_node = 6
        self.root = self.build()
        self.leaves: dict[int, list[int]] = {}
        self.fill_leaves(self.root)

    def find(self, node: int) -> int:
        while self.dsu_parent[node] != node:
            node = self.dsu_parent[node]
        return node

    def build(self) -> int:
        for u, v, weight in EDGES:
            a = self.find(u)
            b = self.find(v)
            if a == b:
                self.logs.append(f"边 ({u},{v},w={weight})：两端已在事件 {a} 下，跳过。")
                continue
            event = self.next_node
            self.next_node += 1
            self.children[event] = [a, b]
            self.value[event] = weight
            self.tree_parent[a] = event
            self.tree_parent[b] = event
            self.tree_parent[event] = None
            self.dsu_parent[a] = event
            self.dsu_parent[b] = event
            self.dsu_parent[event] = event
            self.logs.append(
                f"边 ({u},{v},w={weight})：合并根 {a} 与 {b}，"
                f"新建事件 {event}@{weight}。"
            )
        return self.find(1)

    def fill_leaves(self, node: int) -> list[int]:
        if not self.children[node]:
            answer = [node]
        else:
            answer = []
            for child in self.children[node]:
                answer.extend(self.fill_leaves(child))
        self.leaves[node] = sorted(answer)
        return self.leaves[node]

    def ancestors(self, node: int) -> list[int]:
        chain = [node]
        while self.tree_parent[chain[-1]] is not None:
            chain.append(self.tree_parent[chain[-1]])
        return chain

    def lca(self, u: int, v: int) -> int:
        seen = set(self.ancestors(u))
        for node in self.ancestors(v):
            if node in seen:
                return node
        raise RuntimeError("同一棵重构树中一定有 LCA")

    def component_at(self, vertex: int, threshold: int) -> int:
        node = vertex
        while (
            self.tree_parent[node] is not None
            and self.value[self.tree_parent[node]] is not None
            and self.value[self.tree_parent[node]] <= threshold
        ):
            node = self.tree_parent[node]
        return node

    def print_tree(self, node: int | None = None, prefix: str = "", side: str = "根") -> None:
        if node is None:
            node = self.root
        children = self.children[node]
        if len(children) == 2:
            self.print_tree(children[1], prefix + "    ", "右")
        label = str(node) if self.value[node] is None else f"{node}@{self.value[node]}"
        print(f"{prefix}{side}: {label}  覆盖叶子 {self.leaves[node]}")
        if len(children) == 2:
            self.print_tree(children[0], prefix + "    ", "左")

    def show(self) -> None:
        print("\n成功合并的时间线：")
        for line in self.logs:
            print(" ", line)
        print("\n重构树（@ 后为发生权值）：")
        self.print_tree()
        print("例如 LCA(2,4)=", self.lca(2, 4), "，权值=", self.value[self.lca(2, 4)])


def warmup() -> None:
    print("=== 任务 0：先让并查集健忘一次 ===")
    print("边按权为 2、3、5、6、9 处理。")
    try:
        input("请预测：2 与 4 首次同一泡泡是在权值几？按回车： ")
    except EOFError:
        pass
    print("准备核对。注意：最终 DSU 只告诉你“五个人在一起”，不会保存这个答案。")


def main() -> None:
    tree = ReconstructionTree()
    print("Kruskal 重构树小实验室：事件节点 6..9 是每次成功 union 的合影。")
    warmup()
    tree.show()
    print("\n命令：show；q 2 4；at 2 3；at 2 6；q。")
    while True:
        try:
            raw = input("\n> ").strip()
        except EOFError:
            raw = "q"
        if not raw or raw in {"q", "quit", "exit"}:
            return
        parts = raw.split()
        if parts == ["show"]:
            tree.show()
        elif len(parts) == 3 and parts[0] == "q":
            try:
                u, v = int(parts[1]), int(parts[2])
                if not (1 <= u <= 5 and 1 <= v <= 5):
                    raise ValueError
                event = tree.lca(u, v)
                print(
                    f"{u} 与 {v} 的最近共同事件是 {event}@{tree.value[event]}；"
                    f"首次连通阈值为 {tree.value[event]}。"
                )
            except ValueError:
                print("点编号应在 1..5。")
        elif len(parts) == 3 and parts[0] == "at":
            try:
                vertex, threshold = int(parts[1]), int(parts[2])
                if not 1 <= vertex <= 5:
                    raise ValueError
                event = tree.component_at(vertex, threshold)
                print(
                    f"阈值 {threshold} 时，{vertex} 上升到事件 {event}"
                    f"；所在块的原叶为 {tree.leaves[event]}。"
                )
            except ValueError:
                print("格式示例：at 2 5。")
        else:
            print("小提示：q 2 4 问首次连通；at 2 3 问某阈值时的块。")


if __name__ == "__main__":
    main()