【学习笔记】树状数组
树状数组是一种数据结构，普通树状数组维护的信息及运算要满足结合律且可差分。
一维树状数组
单点加、区间求和
树状数组是用长度为 \(n\) 的数组存储的。我们假设这个数组为 \(n\)，令 lowbit(i)=i&(-i)，则 \(c_i\) 保存的是向前 lowbit(i) 长度的 \(a\) 数组区间和。
单点加：从 \(i\) 开始，修改所有包含 \(a_i\) 的 \(c_i\)： \(c_i,c_{j=i+lowbit(i)} ,c_{j'^=j+lowbit(j)}\)……
区间求和 \([1,x]\)：累加 \(c_x,c_{x'=x-lowbit(x)} ,c_{x''=x'-lowbit(x')}\)……整个过程就相当于不断把 \(x\) 的最后一个 \(1\) 消掉再不断累加。
区间求和 \([l,r]\)：就是利用差分的思想——\([1,r]-[1,l-1]\)
时间复杂度均为 \(O(\log n)\)。
PS：建树的不同方式
-
直接一个一个修改，时间复杂度为 \(O(n\log n)\)。
-
可以先处理一个前缀和数组 \(s\)，由于 \(c_i\) 表示的区间是 \([i-lowbit(i)+1,i]\)，所以可以用 \(s_i-s_{i-lowbit(i)}\) 来表示 \(c_i\)，时间复杂度为 \(O(n)\)。
区间加、单点查询：
换一种思路，用树状数组 \(c\) 的 \(c_i\) 存储向前 lowbit(i) 长度的 \(a\) 数组差分和，那么区间加、单点查询就转变成了原来的单点修改、区间查询。
例题
CF383C Propagating tree
有一颗 \(n\) 个节点的树，需要支持下列两种操作：
- 把某个点的权值加 \(x\)，其孩子权值加 \(-x\)，其孩子的孩子权值加 \(-(-x)\)……
- 查询某个节点的值
\(1\le n\le m\le 2\times 10^5\)
首先把树形转化为区间，可以想到 dfs 序，而看到区间修改、单点查询，考虑用树状数组维护差分数组。如果改变的是深度为奇数的点，就将差分数组加 \(v\)，否则就减 \(v\)，不难发现，树状数组维护的就相当于维护奇数层和偶数层的差。
而查询的时候也分深度奇偶性，如果是奇数就加上 \(1-x\) 的差分和，否则减去即可。
#include <bits/stdc++.h>
#define lowbit(i) i & (-i)
using namespace std;
const int N = 2e5 + 5;
int n, m, ts;
int a[N], l[N], r[N], c[N], d[N];
vector <int> p[N];
void dfs(int k, int fa) {
d[k] = d[fa] + 1;
l[k] = ++ts;
for (auto i : p[k])
if (i != fa) dfs(i, k);
r[k] = ts;
}
void add(int x, int y) {
for (int i = x; i <= n; i += lowbit(i))
c[i] += y;
}
int query(int x) {
int res = 0;
for (int i = x; i>= 1; i -= lowbit(i))
res += c[i];
return res;
}
int main() {
cin >> n >> m;
for (int i = 1; i <= n; i++)
cin >> a[i];
for (int i = 1; i < n; i++) {
int u, v;
cin >> u >> v;
p[u].push_back(v);
p[v].push_back(u);
}
dfs(1, 0);
while (m--) {
int op, x, y;
cin >> op >> x;
if (op == 1) {
cin >> y;
if (d[x] & 1) add(l[x], y), add(r[x] + 1, -y);
else add(l[x], -y), add(r[x] + 1, y);
}
else {
int t = query(l[x]);
if (d[x] & 1) cout << a[x] + t;
else cout << a[x] - t;
cout << "\n";
}
}
return 0;
}
区间加、区间求和
设 a 的差分数组为 d，则：
接下来可以用两个树状数组分别维护 \(d_i\) 与 \(d_i\times i\)。若 \(a[l…r]\) 进行区间加 \(v\)，那么维护 \(d_i\) 的树状数组就做 \(l\) 单点加 \(v\)，对 \(r+1\) 单点加 \(-v\)；对于维护 \(d_i\times i\) 的树状数组，则对 \(l\) 单点加 \(l\times v\)，对 \(r+1\) 单点加 \(-v×(r+1)\) 即可。
int a[N], c1[N], c2[N];
void add(int x, int y) {
for (int i = x; i <= n; i += lowbit(i))
c1[i] += y, c2[i] += y * x;
}
int query(int x, int *c) {
int res = 0;
for (int i = x; i >= 1; i -= lowbit(i))
res += c[i];
return res;
}
void updata(int x, int y, int v) {
add(x, v);
add(y + 1, -v);
}
int ret(int x, int y) {
return query(y, c1) * (y + 1) - query(y, c2) - query(x - 1, c1) * x + query(x - 1, c2);
}
二维树状数组
上面的所有内容都是关于一维树状数组的，其修改和查询的时间复杂度均为 \(O(\log n)\)，空间复杂度均为 \(O(n)\)。接下的内容就是“二维树状数组”的，它的实现和一维有点类似。
单点加、子矩阵求和：
令 \(c_{i,j}\) 表示以 \(a_{i,j}\) 为右下角，高 lowbit(i)，宽 lowbit(j) 的子矩阵的和。理解一下，这其实就是树状数组中套了一个树状数组，这一点具体在下文会体现。
单点加
要修改 \(a_{i,j}\)，令 \(i'=i+lowbit(i),i''=i'+lowbit(i')...,j''=j+lowbit(j),j'=j'+lowbit(j')...\)，则修改 \(c_{{i,i',i''...},{j,j',j''...}}\)。结合上文所说的树状数组套树状数组，就是对于所有需要修改的 \(i\)（和普通树状数组一样求法），再分别修改其对应的 \(j\)，就是维护行的树状数组中修改列的树状数组。
查询子矩阵和
求以 \({i_1,j_1}\) 为左上角，\({i_2,j_2}\) 为右下角的子矩阵的和，利用二维差分的思想：\(s_{i_2,j_2 } -s_{i_2,j_1-1} -s_{i_1-1,j_2 } +s_{i_1-1,i_2-1}\)。
\(s_{i,j}\) 表示的是以 \({1,1}\) 为左上角，\({i,j}\) 为右下角的子矩阵的和，令 \(i'=i-lowbit(i),i''=i'-lowbit(i')…j'=j-lowbit(j),j''=j'-lowbit(j')…\)，则 \(s_{i,j} =c_{i,j} +c_{i,j'} +…+c_{i',j} +…\)。不难发现，这其实和普通树状数组的查询也有点像。
int c[N][N][C], a[N][N];
void add(int x, int y, int z, int v) {
for (int i = x; i <= n; i += lowbit(i))
for (int j = y; j <= m; j += lowbit(j))
c[i][j][v] += z;
}
int query(int x, int y, int v) {
int res = 0;
for (int i = x; i >= 1; i -= lowbit(i))
for (int j = y; j >= 1; j -= lowbit(j))
res += c[i][j][v];
return res;
}
int ret(int x, int y, int x1, int y1, int v) {
return query(x1, y1, v) + query(x - 1, y - 1, v) - query(x - 1, y1, v) - query(x1, y - 1, v);
}
子矩阵加、子矩阵求和
思路和区间加、区间求和有些类似，都是运用了差分，只不过是扩展成了二维。
二维上的差分数组：\(d_{i,j}=a_{i,j}-a_{i,j-1}-a_{i-1,j}+a_{i-1,j-1}\)，对于将左上角 \((i_1,j_1)\)、右下角 \((i_2,j_2)\) 的子矩阵加上 \(v\)，就是将 \(d_{i_1,j_1 }\) 和 \(d_{i_2+1,j_2+1}\) 加 \(v\)，\(d_{i_1,j_2+1}\)、\(d_{i_2+1,j_1 }\) 加 \(-v\)。接着推导：
维护 \(d_{x,y},d_{x,y}\times x,d_{x,y}\times y,d_{x,y}\times xy\) 的树状数组，其它的细节比如要修改哪些、怎么查询都和单点加、子矩阵查询的树状数组一样。
顺便说一句，如果是只需要子矩阵修改、单点查询的话，那么就和区间修改、单点查询一样，只需维护二维差分数组再求二维前缀和即可。
int c1[N][N], c2[N][N], c3[N][N], c4[N][N];
void add(int x, int y, int v) {
for (int i = x; i <= n; i += lowbit(i))
for (int j = y; j <= m; j += lowbit(j)) {
c1[i][j] += v;
c2[i][j] += x * v;
c3[i][j] += y * v;
c4[i][j] += x * y * v;
}
}
int query(int x, int y) {
int res = 0;
for (int i = x; i >= 1; i -= lowbit(i))
for (int j = y; j >= 1; j -= lowbit(j))
res += c1[i][j] * (x + 1) * (y + 1) - c2[i][j] * (y + 1) - c3[i][j] * (x + 1) + c4[i][j];
return res;
}
void updata(int i1, int j1, int i2, int j2, int v) {
add(i1, j1, v), add(i2 + 1, j2 + 1, v);
add(i1, j2 + 1, -v), add(i2 + 1, j1, -v);
}
int getsum(int i1, int j1, int i2, int j2) {
return query(i2, j2) - query(i2, j1 - 1) - query(i1 - 1, j2) + query(i1 - 1, j1 - 1);
}
以上是二维树状数组的全部内容，依旧可以发现，无论是子矩阵修改还是单点修改，它们的修改、查询的时间复杂度均为 \(O(\log^2 n)\)，空间复杂度为 \(O(n^2)\)。
其它应用
权值树状数组
权值树状数组的实现和普通树状数组差不多，只不过其维护的信息是权值罢了，其经常用于解决和偏序相关的问题。
例题
洛谷 P1637 三元上升子序列
- 求数列中满足 \(i<j<k\) 且 \(a_i<a_j<a_k\) 的三元组 \((a_i,a_j,a_k )\) 的个数。
- \(n\le3×10^4,a_i\le10^5\)。
枚举中间的元素 \(a_j\)，分别求其左边满足条件的 \(a_i\) 的个数 \(left\) 和其右边满足条件的 \(a_k\) 的个数 \(right\)，则对答案的贡献为 \(left\times right\)。
以 \(left\) 为例，在枚举 \(a_j\) 的同时用权值树状数组维护 \(a_1-a_{j-1}\) 各个权值的个数，\(left\) 则为其中小于 \(a_i\) 的数个数（用权值树状数组查询）。\(right\) 同理。
#include <bits/stdc++.h>
#define int long long
#define lowbit(x) x & (-x)
using namespace std;
const int N = 5e5 + 5;
int n;
int a[N], b[N], h[N];
int t[N];
void updata(int x, int y){
for (int i = x; i <= 1e5; i += lowbit(i))
t[i] += y;
}
int query(int x) {
int res = 0;
for (int i = x; i; i -= lowbit(i))
res = res + t[i];
return res;
}
signed main() {
int m;
cin >> n;
for (int i = 1; i <= n; i++)
cin >> a[i];
int s1[N];
for (int i = 1; i <= n; i++) {
s1[i] = query(a[i] - 1);
updata(a[i], 1);
}
int s2[N];
memset(t, 0, sizeof(t));
for (int i = n; i >= 1; i--) {
s2[i] = n - i - query(a[i]);
updata(a[i], 1);
}
int ans = 0;
for (int i = 1; i <= n; i++)
ans += s1[i] * s2[i];
cout << ans << "\n";
return 0;
}
树状数组套权值线段树
树状数组套权值线段树就是把权值线段树看作一个整体，把普通线段树中维护的值换作一颗颗线段树而已。
例题
洛谷 P2617 Dynamic Rankings
单点修改、查询区间第 \(k\) 小
\(1\le n,m\le 10^5\)
如果是静态查询区间第 \(k\) 小可以使用可持久化线段树（主席树），但如果是动态的，修改时它的复杂度就会变成 \(O(n\log n)\)，可以用树状数组套权值线段树来优化：主席树的思想为维护 \([1,1][1,2]...[1,n]\) 的权值线段树，其管理的只是一个区间，而树状数组套权值线段树就不一样了，\([1,r]\) 管理的是 \([1,r-lowbit(r)+1]-[1,r]\) 的权值线段树。注意，这里的树状数组并没有一个具体的数组 \(c\)，而只是借用了其的思想而已。
单点查询：在树状数组上求出要修改哪几颗线段树，直接递归修改，这些操作和普通线段树的修改操作相同。
查询区间第 \(k\) 小：由于所有线段树的结构是相同的，所以可以直接在多棵树上进行线段树上二分，其他和用线段树查询差不多，只不过是先用树状数组上记录所有需要的线段树二分用。
单次操作时间复杂度：\(O(\log^2 n)\)
空间复杂度（动态开点）：\(O(n\log^2 n)\)
以下是核心代码（省略主函数的输入、离散化、处理询问）：
const int N = 1e5 + 5;
struct node {
int l, r;
int v;
}t[N * 200];
int n, m, cnt, len, ts;
int a[N], b[N * 2], c[N], l1[N], l2[N], l3[N];
int rt[N];
char op[N];
void pushup(int k) {
t[k].v = t[t[k].l].v + t[t[k].r].v;
}
int updata(int k, int l, int r, int x, int v) {
if (!k) k = ++ts;
if (l == r) {
t[k].v += v;
return k;
}
int m = l + ((r - l) >> 1);
if (x <= m) t[k].l = updata(t[k].l, l, m, x, v);
else t[k].r = updata(t[k].r, m + 1, r, x, v);
pushup(k);
return k;
}
int query(int t1[], int t2[], int s1, int s2, int l, int r, int k) {
if (l == r) return l;
int m = l + ((r - l) >> 1), ls = 0;
for (int i = 1; i <= s1; i++)
ls += t[t[t1[i]].l].v;
for (int i = 1; i <= s2; i++)
ls -= t[t[t2[i]].l].v;
if (ls >= k) {
for (int i = 1; i <= s1; i++) t1[i] = t[t1[i]].l;
for (int i = 1; i <= s2; i++) t2[i] = t[t2[i]].l;
return query(t1, t2, s1, s2, l, m, k);
}
for (int i = 1; i <= s1; i++) t1[i] = t[t1[i]].r;
for (int i = 1; i <= s2; i++) t2[i] = t[t2[i]].r;
return query(t1, t2, s1, s2, m + 1, r, k - ls);
}
void add(int x, int v) {
int xx = lower_bound(b + 1, b + 1 + len, a[x]) - b;
for (int i = x; i <= n; i += lowbit(i))
rt[i] = updata(rt[i], 1, len, xx, v);
}
int init_query(int l, int r, int k) {
int t1[N], t2[N], s1 = 0, s2 = 0;
for (int i = r; i >= 1; i -= lowbit(i))
t1[++s1] = rt[i];
for (int i = l - 1; i >= 1; i -= lowbit(i))
t2[++s2] = rt[i];
return query(t1, t2, s1, s2, 1, len, k);
}
洛谷 P3810 【模板】三维偏序（陌上花开）
有 $ n $ 个元素，第 $ i $ 个元素有 $ a_i,b_i,c_i $ 三个属性，设 $ f(i) $ 表示满足 $ a_j \leq a_i $ 且 $ b_j \leq b_i $ 且 $ c_j \leq c_i $ 且 $ j \ne i $ 的 \(j\) 的数量。对于 $ d \in [0, n) $，求 $ f(i) = d $ 的数量。
\(1\le n\le 10^5\)
这道题可以用许多高级的算法完成，这边介绍一种用树状数组套权值线段树的做法。
从简到繁考虑，一维偏序直接排个序就可以了，二维偏序可以以 \(a_i\) 为关键字排个序，然后用权值树状数组/权值线段树查询即可。而三维偏序也是同样的思路，先以 \(a_i\) 为关键字排序，然后查询 \(b_j\le b_i\) 且 \(c_j\le c_i\) 的个数（\(j\ne i\)）。
不难想到用二维树状数组，行维护 \(b\)，列维护 \(c\)，修改是单点修改，查询则是查询左上角为 \((1,1)\)，右下角为 \((b_i,c_i)\) 的子矩形和。但很可惜，这样的空间复杂度（离散化）是 \(O(n^2)\)，不可实现。
那么可以把其中一个树状数组换换为线段树，在树状数组中维护 \(b\)，在线段树种维护 \(c\)，查询时先在树状数组中查找维护所有 \(\le b_i\) 的线段树，再在这些线段树中查找所有 \(\le c_i\) 的个数即可。
由于这道题是包括等于的，所以应先把所有相等的 \(a_i\) 加入“树状数组”再进行查询，注意这样得到的 \(f(i)\) 会多 \(1\)（包括自己）。
#include <bits/stdc++.h>
#define lowbit(i) i & (-i)
#define endl "\n"
using namespace std;
const int N = 1e5 + 5;
struct getcin {
int x, y, z;
int id;
}a[N];
struct node {
int l, r;
int v;
}t[N << 7];
int n, kkk, ts, cnt;
int tmp[N << 2], rt[N << 2], f[N], ans[N];
bool cmp(getcin x, getcin y) {
if (x.x != y.x) return x.x < y.x;
if (x.y != y.y) return x.y < y.y;
return x.z < y.z;
}
void pushup(int k) {
t[k].v = t[t[k].l].v + t[t[k].r].v;
}
int updata(int k, int l, int r, int x, int v) {
if (!k) k = ++ts;
if (l == r) {
t[k].v += v;
return k;
}
int m = l + ((r - l) >> 1);
if (x <= m) t[k].l = updata(t[k].l, l, m, x, v);
else t[k].r = updata(t[k].r, m + 1, r, x, v);
pushup(k);
return k;
}
int query(int l, int r, int x) {
if (l == r) {
int res = 0;
for (int i = 1; i <= cnt; i++)
res += t[tmp[i]].v;
return res;
}
int m = l + ((r - l) >> 1);
if (x <= m) {
for (int i = 1; i <= cnt; i++)
tmp[i] = t[tmp[i]].l;
return query(l, m, x);
}
int sum = 0;
for (int i = 1; i <= cnt; i++) {
sum += t[t[tmp[i]].l].v;
tmp[i] = t[tmp[i]].r;
}
return sum + query(m + 1, r, x);
}
void add(int x, int y, int v) {
for (int i = x; i <= kkk; i += lowbit(i))
rt[i] = updata(rt[i], 1, kkk, y, v);
}
int find(int x, int y) {
cnt = 0;
for (int i = x; i >= 1; i -= lowbit(i))
tmp[++cnt] = rt[i];
return query(1, kkk, y);
}
int main(){
cin >> n >> kkk;
for (int i = 1; i <= kkk; i++)
rt[i] = ++ts;
for (int i = 1; i <= n; i++)
cin >> a[i].x >> a[i].y >> a[i].z, a[i].id = i;
sort(a + 1, a + 1 + n, cmp);
for (int i = 1; i <= n;) {
int j = i;
do {
add(a[j].y, a[j].z, 1);
j++;
}while (a[j].x == a[j - 1].x);
j = i;
do {
f[a[j].id] = find(a[j].y, a[j].z) - 1;
j++;
}while (a[j].x == a[j - 1].x);
i = j;
}
for (int i = 1; i <= n; i++)
ans[f[i]]++;
for (int i = 0; i < n; i++)
cout << ans[i] << endl;
return 0;
}
练习题
- CF1042D Petya and Array
- P2184 贪婪大陆
- P3586 LOG
- CF961E Tufurama
- CF896E The Untended Antiquity
- P3380 【模板】二逼平衡树（树套树）
参考资料：OI Wiki
