线段树&树状数组

3,290 阅读2分钟

线段树:高效处理「区间和」查询(不仅仅可以处理和、还可以处理区间最值等);二叉树;单点修改。

树状数组:高效处理「前缀和」查询;多叉树;单点修改。

线段树和树状数组不能处理输入数组的长度有增加或者减少的情况。

线段树

参考资料: LeetCode 307. 区域和检索 - 数组可修改 官方题解

线段树是一种非常灵活的数据结构,它可以用于解决多种范围查询问题,比如在对数时间内从数组中找到最小值、最大值、总和、最大公约数、最小公倍数等。

采用数组来构造和储存线段树,设原始数组numsnums长度为nn,则其对应数组实现的线段树长度为2n2n

如果索引 ii 处的元素不是一个叶节点,那么其左子节点和右子节点分别存储在索引为 2i2i2i+12i+1 的元素处。该特性要求线段树不使用数组中下标为0的元素。该规律对任意长度的原始数组都成立。

当数组长度为2n2^n时,对应的线段树是一棵满二叉树。

当数组长度不为2n2^n时,对应的线段树是一棵满二叉树。

线段树可以分为以下三个步骤:

  1. 从给定数组构建线段树的预处理步骤。
  2. 修改元素时更新线段树。
  3. 使用线段树进行区域和检索。
class NumArray:

    def __init__(self, nums: List[int]):
        # 线段树
        if len(nums):
            self.n = len(nums)
            # 线段树的结点数不超过2*n
            self.tree = [0]*2*self.n
            j = 0
            # 先初始化叶子节点, 叶结点放在序号n以后
            for i in range(self.n, 2*self.n):
                self.tree[i] = nums[j]
                j += 1
            # 汇总结点的左子结点下标为2*i, 右子结点下标为2*i+1
            for i in range(self.n-1, 0, -1):
                self.tree[i] = self.tree[2*i]+self.tree[2*i+1]


    def update(self, i: int, val: int) -> None:
        pos = i+self.n
        delta = val-self.tree[pos]
        self.tree[pos] += delta
        while pos:
            pos = pos//2
            self.tree[pos] += delta


    def sumRange(self, i: int, j: int) -> int:
        left = i + self.n
        right = j + self.n
        ans = 0
        while left<=right:
            # 如果左结点指针为奇数,说明当前左结点指针指向区间的右结点,
            # 区间[i,j]不完全包含当前小区间
            if left%2==1:
                # 把当前小区间内属于区间[i,j]的值加到区间和
                ans += self.tree[left]
                # 左结点指针指定右侧完全属于区间[i,j]的区间
                left += 1
            if right%2==0:
                ans += self.tree[right]
                right -= 1
            # 移到上一层
            left = left//2
            right = right//2
        return ans

树状数组

参考资料:315. 计算右侧小于当前元素的个数 题解

树状数组用于计算数组的前缀和。

构造的树状数组不含原始数组numsnums,树状数组长度为原始数组的长度nn+1。

和线段树相同,树状数组不使用索引为0的位置。

树状数组中C[i]C[i]涵盖哪些数是由其索引为ii决定的。把下标写成二进制的形式,最低位的 1 以及后面的 0 表示了预处理数组 C 管理了多少输入数组 A 的元素

这样组织数据,从叶子结点到父结点是可以通过一个叫做 lowbit 的函数计算出来,并且可以知道小于等于当前下标的同一层结点的所有结点

lowbit(x) = x & (-x)

正数的相反数的二进制表示就等于对这个正数的二进制取反+1。以二进制数11010为例,11010的补码为00101,加1后为00110,两者相与便是最低位的1。其实很好理解,补码和原码必然相反,所以原码有0的部位补码全是1,补码再+1之后由于进位那么最末尾的1和原码。

和线段树相同,树状数组可以分为以下三个步骤:

  1. 从给定数组构建树状数组的预处理步骤。
  2. 修改元素时更新树状数组。
  3. 使用树状数组进行前缀和检索。
class FenwickTree:
    def __init__(self, nums):
        self.n = len(nums)
        # 初始化树状数组
        self.tree = [0]*(len(nums)+1)
        for i in range(len(nums)):
            self.update(idx+1, nums[i])

    # 取x的二进制下最右1个1与之后的0组成的数大小
    def _low_bit(self,x):
        return x&(-x)

    # 单点更新:从下到上,最多到 size,可以取等
    def update(self, idx, delta):
        while idx<=self.n:
            self.tree[idx] += delta
            idx += self._low_bit(idx)
    
    # 区间查询:从上到下,最少到 1,可以取等
    def query(self,idx):
        ans = 0
        while idx>=0:
            ans += self.tree[idx]
            idx -= self._low_bit(idx)
        return ans