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