当前位置：首页 > article >正文

树状数组（Binary Indexed Tree/Fenwick Tree）详解

article 2024/12/25 2:12:22

写在前面： 此博客内容已经同步到我的博客网站，如需要获得更优的阅读体验请前往https://blog.mainjay.cloudns.ch/blog/algorithm/binary-indexed-tree

1. 为什么需要树状数组？

1.1 问题起源

考虑这样一个场景：有一个长度为 n 的数组，我们需要：

查询前缀和（即从数组起始位置到某个位置的所有数字之和）
修改数组中某个元素的值

对于这类问题，我们有几种解决方案：

方案1：普通数组

单点修改：O(1)
前缀和查询：O(n)
空间复杂度：O(n)

方案2：前缀和数组

单点修改：O(n)
前缀和查询：O(1)
空间复杂度：O(n)

方案3：树状数组

单点修改：O(log n)
前缀和查询：O(log n)
空间复杂度：O(n)

树状数组在修改和查询操作之间取得了很好的平衡，而且实现简单，常数较小。

1.2 树状数组的本质

树状数组本质上是一个二进制思想的巧妙运用，它通过维护部分和的方式，实现高效的修改和查询操作。每个节点存储的是一个区间和，这个区间的长度是当前位置的二进制中最低位1所表示的值。

2. 树状数组的设计思想

2.1 核心思想

二进制分段：将序列按照二进制拆分，每个节点管理一段区间和
快速统计：利用二进制的性质，实现快速的区间和计算
高效更新：修改某个位置的值时，只需要更新少量节点

2.2 为什么这样设计是高效的？

二进制特性：
- 任何正整数都可以表示为若干个2的幂的和
- 这种特性使得我们可以用O(log n)个节点表示任意区间
区间管理：
- 每个节点管理的区间长度是2的幂
- 这种设计使得更新和查询操作都非常高效

3. 树状数组的应用场景

3.1 实际应用场景

动态排名统计
- 问题：维护一个序列，支持单点修改，查询某个数的排名
- 为什么用树状数组：可以O(log n)时间完成修改和查询
频率统计
- 问题：统计某个元素出现的次数，支持动态修改
- 为什么用树状数组：高效支持更新和查询操作
前缀统计
- 问题：维护前缀和、前缀最大值等信息
- 为什么用树状数组：比线段树实现更简单，常数更小

3.2 算法题中的应用

区间查询类问题
- 例：查询区间和
- 例：查询区间内小于某个值的数的个数
动态维护信息
- 例：动态维护序列的逆序对数量
- 例：动态维护序列的第k大值

4. 基本操作

4.1 完整的树状数组实现

// BIT 树状数组结构
type BIT struct {
    tree []int64
    n    int
}

// NewBIT 初始化树状数组
func NewBIT(n int) *BIT {
    return &BIT{
        tree: make([]int64, n+1), // 下标从1开始
        n:    n,
    }
}

// lowbit 获取x的二进制表示中最低位1所表示的值
func (b *BIT) lowbit(x int) int {
    return x & (-x)
}

// Add 在位置i上增加val
func (b *BIT) Add(i int, val int64) {
    for i <= b.n {
        b.tree[i] += val
        i += b.lowbit(i)
    }
}

// Sum 求前缀和[1,i]
func (b *BIT) Sum(i int) int64 {
    var sum int64
    for i > 0 {
        sum += b.tree[i]
        i -= b.lowbit(i)
    }
    return sum
}

// Query 查询区间和[l,r]
func (b *BIT) Query(l, r int) int64 {
    return b.Sum(r) - b.Sum(l-1)
}

// 使用示例
func Example() {
    // 创建大小为10的树状数组
    bit := NewBIT(10)
    
    // 在位置1添加值5
    bit.Add(1, 5)
    
    // 在位置3添加值3
    bit.Add(3, 3)
    
    // 查询前缀和[1,3]
    sum := bit.Query(1, 3) // 结果为8
}

4.2 二维树状数组

// BIT2D 二维树状数组结构
type BIT2D struct {
    tree [][]int64
    n, m int
}

// NewBIT2D 初始化二维树状数组
func NewBIT2D(n, m int) *BIT2D {
    tree := make([][]int64, n+1)
    for i := range tree {
        tree[i] = make([]int64, m+1)
    }
    return &BIT2D{
        tree: tree,
        n:    n,
        m:    m,
    }
}

// lowbit 获取x的二进制表示中最低位1所表示的值
func (b *BIT2D) lowbit(x int) int {
    return x & (-x)
}

// Add 在位置(x,y)上增加val
func (b *BIT2D) Add(x, y int, val int64) {
    for i := x; i <= b.n; i += b.lowbit(i) {
        for j := y; j <= b.m; j += b.lowbit(j) {
            b.tree[i][j] += val
        }
    }
}

// Sum 求二维前缀和[1,x]×[1,y]
func (b *BIT2D) Sum(x, y int) int64 {
    var sum int64
    for i := x; i > 0; i -= b.lowbit(i) {
        for j := y; j > 0; j -= b.lowbit(j) {
            sum += b.tree[i][j]
        }
    }
    return sum
}

// Query 查询二维区间和[(x1,y1),(x2,y2)]
func (b *BIT2D) Query(x1, y1, x2, y2 int) int64 {
    return b.Sum(x2, y2) - b.Sum(x2, y1-1) - b.Sum(x1-1, y2) + b.Sum(x1-1, y1-1)
}