lessons/互动程序/11-平衡树/00-让刚访问的点靠近你.py

#!/usr/bin/env python3
"""课前零食：一个可查看 zig、zig-zig、zig-zag 的小型 Splay 集合。"""

from __future__ import annotations

from dataclasses import dataclass


@dataclass
class Node:
    key: int
    left: Node | None = None
    right: Node | None = None
    parent: Node | None = None
    size: int = 1


def size(node: Node | None) -> int:
    return node.size if node else 0


class SplaySet:
    def __init__(self) -> None:
        self.root: Node | None = None

    def pull(self, node: Node | None) -> None:
        if node:
            node.size = 1 + size(node.left) + size(node.right)

    def pull_to_root(self, node: Node | None) -> None:
        while node:
            self.pull(node)
            node = node.parent

    def rotate(self, x: Node, log: list[str]) -> None:
        p = x.parent
        if p is None:
            return
        g = p.parent
        if x is p.left:
            middle = x.right
            x.right = p
            p.left = middle
            direction = "右旋"
        else:
            middle = x.left
            x.left = p
            p.right = middle
            direction = "左旋"
        if middle:
            middle.parent = p
        p.parent = x
        x.parent = g
        if g is None:
            self.root = x
        elif p is g.left:
            g.left = x
        else:
            g.right = x
        self.pull(p)
        self.pull(x)
        log.append(f"{direction}：让 {x.key} 绕 {p.key} 上来。")

    def is_right(self, node: Node) -> bool:
        return node.parent is not None and node.parent.right is node

    def splay(self, x: Node, goal: Node | None, log: list[str]) -> None:
        while x.parent is not goal:
            p = x.parent
            if p is None:
                break
            g = p.parent
            if g is goal:
                log.append("zig：父亲已是目标父亲。")
                self.rotate(x, log)
            elif self.is_right(x) == self.is_right(p):
                log.append("zig-zig：祖孙同向，先旋父亲再旋当前点。")
                self.rotate(p, log)
                self.rotate(x, log)
            else:
                log.append("zig-zag：祖孙异向，当前点连续旋两次。")
                self.rotate(x, log)
                self.rotate(x, log)

    def insert_plain(self, key: int) -> None:
        """仅用于造出固定的初始 BST；不触发 Splay。"""
        if self.root is None:
            self.root = Node(key)
            return
        current = self.root
        while True:
            if key < current.key:
                if current.left is None:
                    current.left = Node(key=key, parent=current)
                    self.pull_to_root(current)
                    return
                current = current.left
            else:
                if current.right is None:
                    current.right = Node(key=key, parent=current)
                    self.pull_to_root(current)
                    return
                current = current.right

    def locate(self, key: int) -> tuple[Node | None, Node | None]:
        current = self.root
        last: Node | None = None
        while current:
            last = current
            if key == current.key:
                return current, last
            current = current.left if key < current.key else current.right
        return None, last

    def find(self, key: int) -> tuple[bool, list[str]]:
        node, last = self.locate(key)
        log: list[str] = []
        if node:
            log.append(f"找到 {key}，开始把它伸到根。")
            self.splay(node, None, log)
            return True, log
        if last:
            log.append(f"没找到 {key}，但把最后访问的 {last.key} 伸到根。")
            self.splay(last, None, log)
        return False, log

    def add(self, key: int) -> list[str]:
        if self.root is None:
            self.root = Node(key)
            return [f"空树插入 {key}，它就是根。"]
        node, last = self.locate(key)
        log: list[str] = []
        if node:
            log.append(f"{key} 已存在，仍把它伸到根。")
            self.splay(node, None, log)
            return log
        assert last is not None
        fresh = Node(key=key, parent=last)
        if key < last.key:
            last.left = fresh
        else:
            last.right = fresh
        self.pull_to_root(last)
        log.append(f"把 {key} 当普通 BST 新叶插入，再 Splay。")
        self.splay(fresh, None, log)
        return log

    def delete(self, key: int) -> tuple[bool, list[str]]:
        found, log = self.find(key)
        if not found:
            return False, log
        assert self.root is not None and self.root.key == key
        left, right = self.root.left, self.root.right
        if left:
            left.parent = None
        if right:
            right.parent = None
        self.root.left = self.root.right = None
        if left is None:
            self.root = right
            log.append("左树为空，右树直接成为新根。")
            return True, log
        self.root = left
        maximum = left
        while maximum.right:
            maximum = maximum.right
        log.append(f"左树最大点是 {maximum.key}；在左树内把它伸根后接右树。")
        self.splay(maximum, None, log)
        assert self.root is maximum
        self.root.right = right
        if right:
            right.parent = self.root
        self.pull(self.root)
        return True, log

    def inorder(self) -> list[int]:
        def dfs(node: Node | None) -> list[int]:
            if node is None:
                return []
            return dfs(node.left) + [node.key] + dfs(node.right)
        return dfs(self.root)

    def show_tree(self, node: Node | None = None, prefix: str = "", side: str = "根") -> None:
        if node is None and not prefix:
            node = self.root
        if node is None:
            return
        self.show_tree(node.right, prefix + "    ", "右")
        parent = node.parent.key if node.parent else "-"
        print(f"{prefix}{side}: {node.key} [parent={parent}, size={node.size}]")
        self.show_tree(node.left, prefix + "    ", "左")

    def check(self) -> tuple[bool, str]:
        def dfs(node: Node | None, low: int | None, high: int | None, parent: Node | None) -> tuple[bool, int, str]:
            if node is None:
                return True, 0, ""
            if node.parent is not parent:
                return False, 0, f"{node.key} 的 parent 指针不对"
            if (low is not None and node.key <= low) or (high is not None and node.key >= high):
                return False, 0, f"{node.key} 的 BST 范围不对"
            ok_l, left_size, why = dfs(node.left, low, node.key, node)
            if not ok_l:
                return False, 0, why
            ok_r, right_size, why = dfs(node.right, node.key, high, node)
            if not ok_r:
                return False, 0, why
            if node.size != left_size + 1 + right_size:
                return False, 0, f"{node.key} 的 size 应为 {left_size + 1 + right_size}"
            return True, node.size, ""
        ok, _, why = dfs(self.root, None, None, None)
        return ok, why or "BST、parent、size 均正确"


def print_log(log: list[str]) -> None:
    for line in log:
        print(" ", line)


def warmup() -> None:
    print("=== 任务 0：先预测 7 怎样上来 ===")
    print("初始树根是 10，5 在左，7 是 5 的右儿子。")
    try:
        input("请写下它是 zig-zig 还是 zig-zag，再按回车： ")
    except EOFError:
        pass
    print("核对：7 相对 5 向右，5 相对 10 向左，是 zig-zag。")


def main() -> None:
    tree = SplaySet()
    for key in (10, 5, 15, 3, 7, 12, 18):
        tree.insert_plain(key)
    print("Splay 基本操作小实验室：固定初始 BST，不使用随机数。")
    warmup()
    tree.show_tree()
    print("中序：", tree.inorder())
    print("\n命令：find 7；add 6；del 10；show；check；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()
        try:
            if parts == ["show"]:
                tree.show_tree()
                print("中序：", tree.inorder())
            elif parts == ["check"]:
                ok, message = tree.check()
                print("通过：" if ok else "失败：", message)
            elif len(parts) == 2 and parts[0] == "find":
                found, log = tree.find(int(parts[1]))
                print_log(log)
                print("找到。" if found else "未找到。", "当前根：", tree.root.key if tree.root else "空")
            elif len(parts) == 2 and parts[0] == "add":
                print_log(tree.add(int(parts[1])))
                print("当前根：", tree.root.key if tree.root else "空")
            elif len(parts) == 2 and parts[0] == "del":
                deleted, log = tree.delete(int(parts[1]))
                print_log(log)
                print("已删除。" if deleted else "没有该键。")
            else:
                raise ValueError
        except ValueError:
            print("小提示：命令示例是 find 7、add 6、del 10。")


if __name__ == "__main__":
    main()