#!/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()