lessons/互动程序/08-高级线段树与树状数组/00-会变化的区间和.py

#!/usr/bin/env python3
"""课前零食：在同一数组上观察 Fenwick 与线段树的区间摘要。"""

from __future__ import annotations


DEFAULT = [3, 1, 4, 1, 5, 9]


def ask(prompt: str, default: str = "q") -> str:
    try:
        return input(prompt).strip() or default
    except EOFError:
        return default


def lowbit(value: int) -> int:
    return value & -value


class Fenwick:
    def __init__(self, values: list[int]) -> None:
        self.n = len(values)
        self.tree = [0] * (self.n + 1)
        for index, value in enumerate(values, 1):
            self.add(index, value, quiet=True)

    def add(self, index: int, delta: int, quiet: bool = False) -> None:
        touched: list[int] = []
        while index <= self.n:
            self.tree[index] += delta
            touched.append(index)
            index += lowbit(index)
        if not quiet:
            print("Fenwick 更新的摘要下标：", touched)

    def prefix(self, index: int) -> int:
        answer = 0
        blocks: list[tuple[int, int]] = []
        while index > 0:
            answer += self.tree[index]
            blocks.append((index - lowbit(index) + 1, index))
            index -= lowbit(index)
        print("Fenwick 用到的块：", blocks)
        return answer

    def show(self) -> None:
        for index in range(1, self.n + 1):
            left = index - lowbit(index) + 1
            print(f"tree[{index}]={self.tree[index]:>2}，负责 [{left},{index}]")


class SegmentTree:
    def __init__(self, values: list[int]) -> None:
        self.n = 1
        while self.n < len(values):
            self.n *= 2
        self.tree = [0] * (2 * self.n)
        self.tree[self.n:self.n + len(values)] = values
        for index in range(self.n - 1, 0, -1):
            self.tree[index] = self.tree[index * 2] + self.tree[index * 2 + 1]

    def set(self, index: int, value: int) -> None:
        node = self.n + index - 1
        self.tree[node] = value
        touched = [node]
        node //= 2
        while node:
            self.tree[node] = self.tree[node * 2] + self.tree[node * 2 + 1]
            touched.append(node)
            node //= 2
        print("线段树向上重算的节点：", touched)

    def query(self, left: int, right: int) -> int:
        left += self.n - 1
        right += self.n - 1
        answer = 0
        used: list[int] = []
        while left <= right:
            if left % 2 == 1:
                answer += self.tree[left]
                used.append(left)
                left += 1
            if right % 2 == 0:
                answer += self.tree[right]
                used.append(right)
                right -= 1
            left //= 2
            right //= 2
        print("线段树直接拿到的节点：", used)
        return answer


def main() -> None:
    values = DEFAULT[:]
    fenwick = Fenwick(values)
    segment = SegmentTree(values)
    print("初始数组（下标从 1）：", values)
    print("命令：sum l r；set i x；show；q。直接回车退出。")
    while True:
        command = ask("\n> ")
        if command in {"q", "quit", "exit"}:
            break
        parts = command.split()
        try:
            if parts[0] == "show":
                fenwick.show()
            elif parts[0] == "sum" and len(parts) == 3:
                left, right = map(int, parts[1:])
                if not (1 <= left <= right <= len(values)):
                    raise ValueError("下标应在数组范围内")
                fenwick_answer = fenwick.prefix(right) - fenwick.prefix(left - 1)
                segment_answer = segment.query(left, right)
                print(f"区间 [{left},{right}]：Fenwick={fenwick_answer}，线段树={segment_answer}")
            elif parts[0] == "set" and len(parts) == 3:
                index, value = map(int, parts[1:])
                if not 1 <= index <= len(values):
                    raise ValueError("下标应在数组范围内")
                delta = value - values[index - 1]
                values[index - 1] = value
                fenwick.add(index, delta)
                segment.set(index, value)
                print("新数组：", values)
            else:
                raise ValueError("命令格式不对")
        except (ValueError, IndexError) as error:
            print("小提示：", error)


if __name__ == "__main__":
    main()