MT头条

【学习笔记】树状数组

2023-07-25 22:56 由 zhuoyuxuan 发表于 #其他

树状数组是一种数据结构，普通树状数组维护的信息及运算要满足结合律且可差分。

一维树状数组

单点加、区间求和

树状数组是用长度为 \(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，则：

\[\begin{aligned}a[1...r] & = \sum_{i=1}^r\sum_{j=1}^id_j \\& = \sum_{i=1}^rd_i\times(r-i+1)\\&=\sum_{i=1}^rd_i\times(r+1)-\sum_{i=1}^rd_i\times i\\&=(\sum_{i=1}^rd_i)\times(r+1)-\sum_{i=1}^rd_i\times i\end{aligned} \]

接下来可以用两个树状数组分别维护 \(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\)。接着推导：

\[\begin{aligned}a[1...i][1...j] & =\sum_{x=1}^i\sum_{y=1}^j\sum_{h=1}^x\sum_{k=1}^yd_{h,k}\\&=\sum_{x=1}^i\sum_{y=1}^jd_{x,y}\times(i-x+1)\times(j-y+1)\\&=\sum_{x=1}^i\sum_{y=1}^jd_{x,y}\times(ij-iy+i-xj-xy-x+j-y+1)\\&=(\sum_{x=1}^i\sum_{y=1}^jd_{x,y})\times(ij+i+j+1)-(\sum_{x=1}^i\sum_{y=1}^jd_{x,y}\times x)\times(j+1)-(\sum_{x=1}^i\sum_{y=1}^jd_{x,y}\times y)\times(i+1)+\sum_{x=1}^i\sum_{y=1}^jd_{x,y}\times xy\end{aligned} \]

维护 \(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

热门相关：总裁别再玩了   上神来了   修真界败类   神算大小姐   激情社区