lessons/互动程序/11-平衡树/00-从排好队到会分身的树.py

#!/usr/bin/env python3
"""课前零食：可操作地追踪一棵按值 FHQ Treap。"""

from __future__ import annotations

from dataclasses import dataclass
import random


RNG = random.Random(20260903)


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


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


def pull(node: Node | None, log: list[str]) -> None:
    if node:
        node.size = node_size(node.left) + node.count + node_size(node.right)
        log.append(f"pull({node.key})：size={node.size}")


def split(node: Node | None, key: int, log: list[str]) -> tuple[Node | None, Node | None]:
    """返回所有键 <= key 的树，以及所有键 > key 的树。"""
    if node is None:
        return None, None
    if node.key <= key:
        log.append(f"split({node.key}, {key})：{node.key} 放左半，继续切右子树")
        middle, right = split(node.right, key, log)
        node.right = middle
        pull(node, log)
        return node, right
    log.append(f"split({node.key}, {key})：{node.key} 放右半，继续切左子树")
    left, middle = split(node.left, key, log)
    node.left = middle
    pull(node, log)
    return left, node


def merge(left: Node | None, right: Node | None, log: list[str]) -> Node | None:
    """前提：left 的所有键 <= right 的所有键。"""
    if left is None:
        return right
    if right is None:
        return left
    if left.priority > right.priority:
        log.append(f"merge：{left.key}(p={left.priority}) 优先级更高，做根")
        left.right = merge(left.right, right, log)
        pull(left, log)
        return left
    log.append(f"merge：{right.key}(p={right.priority}) 优先级更高，做根")
    right.left = merge(left, right.left, log)
    pull(right, log)
    return right


def insert(root: Node | None, key: int) -> tuple[Node | None, list[str]]:
    log: list[str] = []
    less, at_least = split(root, key - 1, log)
    equal, greater = split(at_least, key, log)
    if equal:
        equal.count += 1
        pull(equal, log)
        log.append(f"键 {key} 已存在：count 增为 {equal.count}")
    else:
        equal = Node(key=key, priority=RNG.randrange(1, 100))
        log.append(f"新建 {key}，随机优先级为 {equal.priority}")
    return merge(less, merge(equal, greater, log), log), log


def erase_one(root: Node | None, key: int) -> tuple[Node | None, list[str]]:
    log: list[str] = []
    less, at_least = split(root, key - 1, log)
    equal, greater = split(at_least, key, log)
    if equal is None:
        log.append(f"没有键 {key}，原树接回去。")
    elif equal.count > 1:
        equal.count -= 1
        pull(equal, log)
        log.append(f"键 {key} 的 count 降为 {equal.count}")
    else:
        log.append(f"删除唯一的键 {key}，把它的左右孩子合并。")
        equal = merge(equal.left, equal.right, log)
    return merge(less, merge(equal, greater, log), log), log


def rank(root: Node | None, key: int) -> int:
    answer = 1
    while root:
        if key <= root.key:
            root = root.left
        else:
            answer += node_size(root.left) + root.count
            root = root.right
    return answer


def kth(root: Node | None, k: int) -> int | None:
    while root:
        left_count = node_size(root.left)
        if k <= left_count:
            root = root.left
        elif k <= left_count + root.count:
            return root.key
        else:
            k -= left_count + root.count
            root = root.right
    return None


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


def show_tree(root: Node | None, prefix: str = "", side: str = "根") -> None:
    if root is None:
        return
    show_tree(root.right, prefix + "    ", "右")
    print(f"{prefix}{side}: {root.key} [p={root.priority}, cnt={root.count}, size={root.size}]")
    show_tree(root.left, prefix + "    ", "左")


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


def warmup() -> None:
    """让实验从数组/BST 的失败开始，而不是直接倾倒递归日志。"""
    print("\n=== 任务 0：先当人脑数据结构 ===")
    print("集合是 [4, 2, 7, 5, 2]。请先在纸上排成 [2, 2, 4, 5, 7]。")
    try:
        input("若插入 6，有序数组中哪几个元素要右移？写下预测后按回车： ")
    except EOFError:
        pass
    print("答案检查：6 放在 5 与 7 之间，至少要把 7 向右搬一格。")
    print("再想一想递增插入 1,2,3,4 的普通 BST：它会变成链。")
    try:
        input("四张优先级卡 2/35、4/60、5/83、7/20 中，谁应做 Treap 根？按回车核对： ")
    except EOFError:
        pass
    print("答案检查：5 的优先级 83 最大；键值左右关系仍由 key 决定。")
    print("接下来的日志只验证你刚搭出的框架：切范围、改目标块、合并、pull。")


def main() -> None:
    root: Node | None = None
    print("FHQ Treap 小实验室：固定随机种子，方便复盘。")
    warmup()
    for value in (4, 2, 7, 5, 2):
        root, log = insert(root, value)
        print(f"\n插入 {value}：")
        print_log(log)
    print("\n初始树：")
    show_tree(root)
    print("中序（应有序）：", inorder(root))
    print("命令：+ x；- x；rank x；k k；show；q。直接回车退出。")
    while True:
        try:
            raw = input("\n> ").strip()
        except EOFError:
            raw = "q"
        if not raw or raw in {"q", "quit", "exit"}:
            break
        parts = raw.split()
        try:
            if parts[0] == "+" and len(parts) == 2:
                root, log = insert(root, int(parts[1]))
                print_log(log)
            elif parts[0] == "-" and len(parts) == 2:
                root, log = erase_one(root, int(parts[1]))
                print_log(log)
            elif parts[0] == "rank" and len(parts) == 2:
                print(f"rank({parts[1]}) = {rank(root, int(parts[1]))}")
            elif parts[0] == "k" and len(parts) == 2:
                value = kth(root, int(parts[1]))
                print("第 k 小：" if value is not None else "k 超出元素总数。", value if value is not None else "")
            elif parts[0] == "show":
                show_tree(root)
                print("中序：", inorder(root))
            else:
                raise ValueError("命令格式不对")
        except ValueError as error:
            print("小提示：", error)


if __name__ == "__main__":
    main()