lessons/互动程序/11-平衡树/00-把相同的一段装进一个盒子.py

#!/usr/bin/env python3
"""课前零食：用可见段卡模拟 ODT 的 split 与区间操作。"""

from __future__ import annotations

from dataclasses import dataclass


@dataclass
class Segment:
    left: int
    right: int
    value: int

    @property
    def length(self) -> int:
        return self.right - self.left + 1

    def show(self) -> str:
        return f"[{self.left},{self.right},v={self.value}]"


class SegmentLab:
    def __init__(self) -> None:
        self.model = [1, 1, 1, 4, 4, 2, 2, 2, 2]
        self.segments = [
            Segment(1, 3, 1),
            Segment(4, 5, 4),
            Segment(6, 9, 2),
        ]

    @property
    def n(self) -> int:
        return len(self.model)

    def normalize(self) -> None:
        merged: list[Segment] = []
        for segment in self.segments:
            if (
                merged
                and merged[-1].right + 1 == segment.left
                and merged[-1].value == segment.value
            ):
                merged[-1].right = segment.right
            else:
                merged.append(Segment(segment.left, segment.right, segment.value))
        self.segments = merged

    def split(self, position: int, log: list[str]) -> int:
        if not 1 <= position <= self.n + 1:
            raise ValueError("切分位置应在 1..n+1")
        if position == self.n + 1:
            log.append("split(n+1)：这是尾后边界，不新建段。")
            return len(self.segments)
        for index, segment in enumerate(self.segments):
            if segment.left <= position <= segment.right:
                if position == segment.left:
                    log.append(f"split({position})：刚好在 {segment.show()} 左端，什么也不切。")
                    return index
                left_piece = Segment(segment.left, position - 1, segment.value)
                right_piece = Segment(position, segment.right, segment.value)
                self.segments[index : index + 1] = [left_piece, right_piece]
                log.append(
                    f"split({position})：{segment.show()} 变成 "
                    f"{left_piece.show()} 与 {right_piece.show()}。"
                )
                return index + 1
        raise RuntimeError("段卡没有完整覆盖数组")

    def bounds(self, left: int, right: int, log: list[str]) -> tuple[int, int]:
        if not 1 <= left <= right <= self.n:
            raise ValueError("区间应满足 1 ≤ l ≤ r ≤ n")
        end_position = right + 1
        self.split(end_position, log)
        begin = self.split(left, log)
        # C++ set 的 end 迭代器在 split(left) 后仍有效；本实验室为可读性
        # 使用列表，左侧插入可能令整数下标右移，故重新定位同一边界。
        if end_position == self.n + 1:
            end = len(self.segments)
        else:
            end = next(
                index
                for index, segment in enumerate(self.segments)
                if segment.left == end_position
            )
        return begin, end

    def add(self, left: int, right: int, delta: int) -> list[str]:
        log: list[str] = []
        begin, end = self.bounds(left, right, log)
        touched = end - begin
        for index in range(begin, end):
            self.segments[index].value += delta
        for index in range(left - 1, right):
            self.model[index] += delta
        self.normalize()
        log.append(f"区间加触及 {touched} 张完整段卡；规范化后共有 {len(self.segments)} 段。")
        return log

    def assign(self, left: int, right: int, value: int) -> list[str]:
        log: list[str] = []
        begin, end = self.bounds(left, right, log)
        removed = end - begin
        self.segments[begin:end] = [Segment(left, right, value)]
        self.model[left - 1 : right] = [value] * (right - left + 1)
        self.normalize()
        log.append(
            f"赋值删除中央 {removed} 张卡，插入一张 [{left},{right},v={value}]；"
            f"规范化后共有 {len(self.segments)} 段。"
        )
        return log

    def kth(self, left: int, right: int, rank: int) -> tuple[int, list[str]]:
        log: list[str] = []
        begin, end = self.bounds(left, right, log)
        pieces = sorted(
            (self.segments[index].value, self.segments[index].length)
            for index in range(begin, end)
        )
        total = sum(length for _, length in pieces)
        if not 1 <= rank <= total:
            raise ValueError("k 超出区间长度")
        original_rank = rank
        for value, length in pieces:
            if rank <= length:
                log.append(f"按值排序后的段贡献为 {pieces}；第 {original_rank} 小落在值 {value}。")
                return value, log
            rank -= length
        raise RuntimeError("长度账本损坏")

    def power_sum(self, left: int, right: int, exponent: int, modulus: int) -> tuple[int, list[str]]:
        if modulus <= 0:
            raise ValueError("模数必须为正")
        log: list[str] = []
        begin, end = self.bounds(left, right, log)
        answer = 0
        for index in range(begin, end):
            segment = self.segments[index]
            answer = (answer + segment.length * pow(segment.value, exponent, modulus)) % modulus
        log.append(f"幂和只遍历了 {end - begin} 张段卡。")
        return answer, log

    def attack(self) -> list[str]:
        self.model = list(range(1, self.n + 1))
        self.segments = [Segment(index, index, value) for index, value in enumerate(self.model, start=1)]
        log = [
            f"现在每个位置值都不同：共有 {len(self.segments)} 张单点段卡。",
            "接下来 add(1,n,1) 仍需扫过所有段；相邻差不变，无法压缩。",
        ]
        return log

    def show(self) -> None:
        print("段卡：", "  ".join(segment.show() for segment in self.segments))
        print("数组：", self.model)
        print("段数：", len(self.segments))

    def check(self) -> tuple[bool, str]:
        expanded: list[int] = []
        previous_right = 0
        for segment in self.segments:
            if segment.left != previous_right + 1 or segment.left > segment.right:
                return False, "段卡没有按顺序无缝覆盖"
            expanded.extend([segment.value] * segment.length)
            previous_right = segment.right
        if previous_right != self.n:
            return False, "段卡没有覆盖到 n"
        if expanded != self.model:
            return False, f"段卡展开 {expanded} 与数组验算 {self.model} 不同"
        return True, "段卡覆盖、数值与普通数组均一致"


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


def warmup() -> None:
    print("=== 任务 0：先猜边界，不要先背 split ===")
    print("初始段卡是 [1,3,1] [4,5,4] [6,9,2]。")
    try:
        input("预测 split(2) 会把哪张卡拆成哪两张，再按回车： ")
    except EOFError:
        pass
    print("核对：[1,3,1] 变为 [1,1,1] 与 [2,3,1]，数组数值完全不变。")


def main() -> None:
    lab = SegmentLab()
    print("ODT 段卡小实验室：这里用列表展示卡片，方便看清 split。")
    warmup()
    lab.show()
    print("\n命令：show；split 2；add 2 7 3；assign 3 8 9；kth 1 9 5；pow 1 9 2 97；attack；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"]:
                lab.show()
            elif parts == ["attack"]:
                print_log(lab.attack())
                lab.show()
            elif parts == ["check"]:
                ok, message = lab.check()
                print("通过：" if ok else "失败：", message)
            elif len(parts) == 2 and parts[0] == "split":
                log: list[str] = []
                lab.split(int(parts[1]), log)
                print_log(log)
                lab.show()
            elif len(parts) == 4 and parts[0] == "add":
                print_log(lab.add(int(parts[1]), int(parts[2]), int(parts[3])))
                lab.show()
            elif len(parts) == 4 and parts[0] == "assign":
                print_log(lab.assign(int(parts[1]), int(parts[2]), int(parts[3])))
                lab.show()
            elif len(parts) == 4 and parts[0] == "kth":
                answer, log = lab.kth(int(parts[1]), int(parts[2]), int(parts[3]))
                print_log(log)
                print("答案：", answer)
            elif len(parts) == 5 and parts[0] == "pow":
                answer, log = lab.power_sum(int(parts[1]), int(parts[2]), int(parts[3]), int(parts[4]))
                print_log(log)
                print("答案：", answer)
            else:
                raise ValueError("命令格式不对")
        except ValueError as error:
            print("小提示：", error)


if __name__ == "__main__":
    main()