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

// 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
}

// 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)
}