树状数组模板

发表于

template类

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
template <typename T>
struct Fenwick
{
    int n;
    std::vector<T> a;
    Fenwick(int n_ = 0)
    {
        init(n_);
    }
    void init(int n_)
    {
        n = n_;
        a.assign(n, T{});
    }
    void add(int x, const T &v)
    {
        for (int i = x + 1; i <= n; i += i & -i) {
            a[i - 1] = a[i - 1] + v;
        }
    }
    T sum(int x) {
        T ans{};
        for (int i = x; i > 0; i -= i & -i) {
            ans = ans + a[i - 1];
        }
        return ans;
    }
    T rangeSum(int l, int r) {
        return sum(r) - sum(l);
    }
    int select(const T &k) {
        int x = 0; T cur{};
        for (int i = 1 << int(std::log2(n)); i; i /= 2) {
            if (x + i <= n && cur + a[x + i - 1] <= k) {
                x += i; cur = cur + a[x - 1];
            }
        }
        return x;
    }
};

Fenwick 树（树状数组）模板解析与用法

模板结构

这个 Fenwick 结构体是 树状数组（Fenwick Tree，或 Binary Indexed Tree, BIT），它用于 高效维护前缀和，支持单点修改和前缀求和操作。

成员变量

1
2
int n;           
std::vector<T> a;
  • n：数组的大小。
  • a：存储树状数组的 核心数据结构，用 std::vector<T> 维护。

构造函数

1
2
3
Fenwick(int n_ = 0) {
    init(n_);
}
  • 作用：初始化 Fenwick 树，默认为空 (n = 0)。
  • init(n_)：调用初始化函数。

函数解析

1. init(int n_)：初始化

1
2
3
4
void init(int n_) {
    n = n_;
    a.assign(n, T{});
}
  • 作用：初始化 Fenwick 树大小 n，并将所有元素清空（设为 T{}，即类型 T 的默认值）。
  • a.assign(n, T{})：创建大小为 n 的数组，所有元素初始化为 0（或 T{}）。

2. add(int x, const T &v)：单点更新

1
2
3
4
5
void add(int x, const T &v) {
    for (int i = x + 1; i <= n; i += i & -i) {
        a[i - 1] = a[i - 1] + v;
    }
}
  • 作用：在索引 x 位置加上 v，维护树状数组结构。
  • 核心原理：
    • i = x + 1：树状数组通常从 1 开始存储，因此 x 需要 +1（但 a 仍从 0 开始）。
    • i += i & -i：跳转到下一个受影响的节点。
    • a[i - 1] += v：累加 v 到对应位置。

3. sum(int x)：求前缀和

1
2
3
4
5
6
7
T sum(int x) {
    T ans{};
    for (int i = x; i > 0; i -= i & -i) {
        ans = ans + a[i - 1];
    }
    return ans;
}
  • 作用：计算前缀 [0, x-1] 的区间和。
  • 核心原理：
    • i = x：直接从 x 开始往上累加贡献值。
    • i -= i & -i：跳转到前一个影响的节点。
    • ans += a[i - 1]：累加贡献值。

4. rangeSum(int l, int r)：求区间和

1
2
3
T rangeSum(int l, int r) {
    return sum(r) - sum(l);
}
  • 作用：求 [l, r-1] 区间的和。
  • 计算方式：
    • sum(r) 计算 [0, r-1] 的前缀和。
    • sum(l) 计算 [0, l-1] 的前缀和。
    • 区间和 = sum(r) - sum(l)。

5. select(const T &k)：二分查找满足前缀和 ≤ k 的最大索引

1
2
3
4
5
6
7
8
9
10
11
int select(const T &k) {
    int x = 0;
    T cur{};
    for (int i = 1 << int(std::log2(n)); i; i /= 2) {
        if (x + i <= n && cur + a[x + i - 1] <= k) {
            x += i;
            cur = cur + a[x - 1];
        }
    }
    return x;
}
  • 作用：找到最大的 x，使得 sum(x) ≤ k。
  • 核心思路：
    • 使用 倍增思想，从最高位开始尝试加 i，找到最大的 x，使得 sum(x) ≤ k。
    • i = 1 << int(std::log2(n))：找最高位的 1。
    • x + i <= n && cur + a[x + i - 1] <= k：
      • 如果当前累加不会超过 k，就把 x 更新到 x + i。
      • cur = cur + a[x - 1]：更新累加值。

用法示例

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
#include <iostream>
#include <vector>

int main() {
    Fenwick<int> fen(10); // 创建大小为 10 的树状数组

    // 在索引 2, 4, 6 位置增加值
    fen.add(2, 5);
    fen.add(4, 3);
    fen.add(6, 7);

    // 计算前缀和
    std::cout << "sum(6) = " << fen.sum(6) << "\n";  // sum(6) = 5 + 3 + 7 = 15
    std::cout << "sum(3) = " << fen.sum(3) << "\n";  // sum(3) = 5
    std::cout << "rangeSum(2, 6) = " << fen.rangeSum(2, 6) << "\n"; // sum(6) - sum(2)

    // 选择满足前缀和 <= 8 的最大索引
    std::cout << "select(8) = " << fen.select(8) << "\n";  

    return 0;
}

总结

函数 作用
init(n_) 初始化树状数组，大小为 n_
add(x, v) 在索引 x 位置加上 v
sum(x) 计算 [0, x-1] 的前缀和
rangeSum(l, r) 计算 [l, r-1] 的区间和
select(k) 找到最大的 x，使得 sum(x) ≤ k

这个 Fenwick 树适用于：

  1. 动态维护数组前缀和（支持修改）
  2. 高效计算区间和（O(log n) 复杂度）
  3. 二分查找前缀和（select(k)）

适合用于 动态前缀和查询、区间查询、逆序对计算 等场景。