Python递归实现树状数组数据结构

2023-04-16 00:00:00 递归 数据结构 树状

树状数组（Fenwick Tree）是一种数据结构，可以在 $O(logn)$ 的时间复杂度下完成单点修改和前缀求和操作。

Python递归实现树状数组：

class BinaryIndexedTree:
    def __init__(self, n: int):
        self.tree = [0] * (n + 1)

    def __lowbit(self, x: int) -> int:
        return x & -x

    def update(self, i: int, delta: int):
        if i == 0:
            return
        self.tree[i] += delta
        self.update(i - self.__lowbit(i), delta)

    def query(self, i: int) -> int:
        if i == 0:
            return 0
        return self.tree[i] + self.query(i - self.__lowbit(i))

在具体操作中，树状数组的编号从 1 开始，因而在数组查询和更新时都需要进行一次加一操作。

例如，我们可以使用树状数组来求一个数组的前缀和：

arr = [1, 3, 5, 7, 9]
n = len(arr)
BIT = BinaryIndexedTree(n)
for i in range(1, n + 1):
    BIT.update(i, arr[i - 1])
    # 注意：数组 arr 的下标从 0 开始，而 BIT 的下标从 1 开始
for i in range(1, n + 1):
    print(BIT.query(i))

输出结果为：

1
4
9
16
25

这样，我们就成功地使用递归实现了树状数组的数据结构！

相关文章