Python中的树形数据结构: 树状数组

2023-04-11 00:00:00 数组数据结构树状

树状数组是一种可以用来维护序列的前缀和的数据结构，也可以用来解决一些区间统计问题。在 Python 中，我们可以使用列表来实现树状数组。

以下是一个简单的树状数组实现：

class FenwickTree:
    def __init__(self, n):
        self.nums = [0] * (n+1)
        self.size = n

    def lowbit(self, x):
        return x & -x

    def update(self, i, val):
        while i <= self.size:
            self.nums[i] += val
            i += self.lowbit(i)

    def query(self, i):
        res = 0
        while i > 0:
            res += self.nums[i]
            i -= self.lowbit(i)
        return res

其中，nums 表示树状数组，size 表示数组的大小。lowbit 函数用来计算一个数的最低位 1 所对应的值。update 函数用来更新某个位置的数值，同时要更新对应的前缀和。query 函数用来查询某个位置之前的前缀和。

例如，我们可以使用以下代码来使用这个树状数组：

arr = [1, 3, 5, 7, 9, 11]
tree = FenwickTree(len(arr))
for i, num in enumerate(arr):
    tree.update(i+1, num)

print(tree.nums)  # [0, 1, 4, 5, 16, 9, 20]
print(tree.query(3))  # 9
print(tree.query(5))  # 25

# 修改 arr[2] = 4
tree.update(3, 1)
arr[2] = 4

print(tree.nums)  # [0, 1, 4, 6, 16, 9, 20]
print(tree.query(3))  # 6
print(tree.query(5))  # 24

上述代码中，我们首先使用树状数组初始化一个数组 arr，然后使用 query 函数来查询某个位置之前的前缀和。接着，我们修改了 arr 中的一个数，并使用 update 函数来更新树状数组中对应位置的数值。最后，我们再次使用 query 函数来验证更新是否正确。

需要注意的是，树状数组的下标从 1 开始，因此在实际使用时需要将数组下标加一。

相关文章