浅谈权值树状数组及其扩展

5k_sync_closer

2022-01-23 14:51:49

算法·理论

权值树状数组是啥?

这是一种当半个平衡树使的树状数组,值域小不强制在线(两者只需满足其一)时可以用。

平衡树的一系列操作时间复杂度均为 O(\log n)n 是值域)。

(如果值域大,可以把操作离线下来然后离散化)

然鹅要开权值数组,所以空间复杂度 O(n)

怎么搞?

以平衡树板题为例:

维护一个集合,6 个操作:

  1. 插入 x
  2. 删除 x
  3. 查询 x 数的排名
  4. 查询排名为 x 的数
  5. x 的前驱
  6. x 的后继

这个树状数组是当做一个桶来用的。

假设这个树状数组维护的桶叫做 g

插入 x

## 删除 x $x$ 的出现次数 $-1$,即修改 $g_x\gets g_x-1$。 ## 求 x 的排名 不难得出,$\left(\sum\limits_{i=1}^{x-1}g_i\right)+1$ 即为所求。 ## 求排名为 x 的数 设 $x$ 的排名表示为 $f(x)=\left(\sum\limits_{i=1}^{x-1}g_i\right)+1$。 形式化地,求 $\max\limits_{f(k)\leq x}k$,我们把下面的条件移个项: $$ \begin{aligned} \left(\sum\limits_{i=1}^{k-1}g_i\right)+1\leq x\\ \sum\limits_{i=1}^{k-1}g_i\leq x-1 \end{aligned} $$ 看起来像二分,不过复杂度变成 $\log^2n$,不太行。 有一个叫树状数组上倍增的 trick(可以去看[这道题](https://www.luogu.com.cn/problem/P6619)的题解),可以解决这个问题。 设当前位置为 $r$(初值为 $0$),$t=\sum\limits_{i=1}^rg_i$。 不断尝试把 $r$ 往后跳 $2^{\log n},2^{\log n-1},2^{\log n-2},...,2^0$,时刻保证 $\sum\limits_{i=1}^rg_i\leq x-1$,即 $t<x$。 考虑如何更新 $t$。不难发现,$r$ 只会从 $r - \operatorname{lowbit}(r)+1$ 跳到 $r$。 于是我们直接利用树状数组的结构 $c_r=\sum\limits_{i=r-\operatorname{lowbit}(r)+1}^rg_i$,$t\gets t+c_r$ 即可。 $r$ 不能再往后跳,就说明 $r$ 达到满足条件的最大值, 对照上面的公式,$r=k-1$,答案即为 $r+1$。 举个例子:集合内有 $8$ 个数 $1,2,3,4,5,6,7,8$: ![](https://cdn.luogu.com.cn/upload/image_hosting/7pssa5lf.png) 现在要求排名为 $6$ 的数。 尝试把 $r$ 向后跳 $2^3$ 到 $r=8$: ![](https://cdn.luogu.com.cn/upload/image_hosting/e20wz6jv.png) $t=8>6$,不能跳,$r=0$,$t=0$。 尝试把 $r$ 向后跳 $2^2$ 到 $r=4$: ![](https://cdn.luogu.com.cn/upload/image_hosting/4zhbao5e.png) $t=4<6$,可以跳,$r=4$,$t=4$。 尝试把 $r$ 向后跳 $2^1$ 到 $r=6$: ![](https://cdn.luogu.com.cn/upload/image_hosting/xumoahc7.png) $t=6=6$,不能跳,$r=4$,$t=4$。 尝试把 $r$ 向后跳 $2^0$ 到 $r=5$: ![](https://cdn.luogu.com.cn/upload/image_hosting/eo1flxul.png) $t=5<6$,可以跳,$r=5$,$t=5$。 答案即为 $r+1=6$。 ```cpp int kth(int k) { int r = 0, t = 0; for(int i = log值域, x, y;i >= 0;--i) { if((x = r + (1 << i)) > 值域) continue; if((y = t + c[x]) < k) r = x, t = y; } return r + 1; } ``` ## 求 x 的前驱 前驱就是排名为 $\max\limits_{f(i)<f(x)}f(i)$ 的数。 要保证 $f(i)<f(x)$ 的情况下 $f(i)$ 最大,可以令 $f(i)=f(x)-1$。 那么只需要求 $\operatorname{kth}(f(x)-1)$ 即可(**注意括号**)。 ## 求 x 的后继 同理,后继就是排名为 $\min\limits_{i>x}f(i)$ 的数。 要保证 $i>x$ 的情况下 $f(i)$ 最小,可以令 $i=x+1$。 即求 $\operatorname{kth}(f(x+1))$ (**注意括号**)。 ------------ 参考代码: ```cpp #include <cmath> #include <cstdio> #include <algorithm> #define h(x) lower_bound(a, a + m, x) - a + 1 using namespace std; struct Q{int o, x;}q[100050]; int n, m, lg, a[100050], c[100050]; void chg(int x, int k) {for(;x <= m;x += x & -x) c[x] += k;} int ask(int x) {int r = 0;for(;x;x &= x - 1) r += c[x];return r;} int kth(int k) { int r = 0, t = 0; for(int i = lg, x, y;i >= 0;--i) { if((x = r + (1 << i)) > m) continue; if((y = t + c[x]) < k) r = x, t = y; } return a[r]; } int main() { scanf("%d", &n); for(int i = 0;i < n;++i) { scanf("%d%d", &q[i].o, &q[i].x); if(q[i].o != 4) a[m++] = q[i].x; } sort(a, a + m);lg = log2(m = unique(a, a + m) - a); for(int i = 0, o, x;i < n;++i) { o = q[i].o, x = q[i].x; switch(o) { case 1: chg(h(x), 1);break; case 2: chg(h(x), -1);break; case 3: printf("%d\n", ask(h(x) - 1) + 1);break; case 4: printf("%d\n", kth(x));break; case 5: printf("%d\n", kth(ask(h(x) - 1)));break; case 6: printf("%d\n", kth(ask(h(x)) + 1));break; } } return 0; } ``` 显而易见,把求和换成任何奇奇怪怪的运算符, **且保证 $f(i)$ 单调不降**(满足倍增性质),就可以维护各种奇奇怪怪的东西。 >重新定义 $x$ 的排名为 $\max\limits_{i=1}^{x-1}c(i)$,其中 $c(i)$ 为 $i$ 在集合中出现的次数。 >在此排名意义下,维护平衡树的 6 个操作。允许 $\log^2n$ 插入删除。值域 $[1,10^6]$。 做法类似,不展开叙述: ```cpp #include <cstdio> #define l(x) (x & -x) int n, a[1000050], c[1000050]; int max(int x, int y) {return x > y ? x : y;} void chg(int x, int k) {a[x] += k;for(;x <= 1e6;x += l(x)) {c[x] = a[x]; for(int i = 1;i < l(x);i <<= 1) c[x] = max(c[x], c[x - i]);}} int ask(int x) {int r = 0;for(;x;x &= x - 1) r = max(r, c[x]);return r;} int kth(int k) { int r = 0, t = 0; for(int i = 20, x, y;i >= 0;--i) { if((x = r + (1 << i)) > 1e6) continue; if((y = max(t, c[x])) <= k) r = x, t = y; } return r + 1; } int main() { scanf("%d", &n); for(int i = 0, o, x;i < n;++i) { scanf("%d%d", &o, &x); switch(o) { case 1: chg(x, 1);break; case 2: chg(x, -1);break; case 3: printf("%d\n", ask(x - 1));break; case 4: printf("%d\n", kth(x));break; case 5: printf("%d\n", kth(ask(x - 1) - 1));break; case 6: printf("%d\n", kth(ask(x)));break; } } return 0; } ``` 注意到权值线段树可以 $\log n$ 插入删除,但 BIT 常数小无伤大雅。 当然,你也可以用类似的方法解决[区间 MEX 问题](https://www.luogu.com.cn/problem/P4137), 但维护 MEX 的树状数组并不是(狭义的)权值树状数组,本文不做说明。 # 扩展篇 >考虑维护一种数据结构: >1. 插入 $i∈[x,y]$。 >2. 删除 $i∈[x,y]$。 >3. 求 $x$ 的排名。 >4. 求排名为 $x$ 的数。 >5. 求 $x$ 的前驱。 >6. 求 $x$ 的后继。 权值线段树可以轻松解决这种问题,但众所周知 BIT 常数更优。 (注意到这里二者都可以 $\log n$ 维护所有操作) 我们知道,树状数组维护区间修区间查用两个数组, 一个维护差分数组 $d_i$,一个维护 $d_i\times (i-1)$。 那么查询前缀和就可以这样: $$ \begin{aligned} \sum\limits_{i=1}^xg_i&=\sum\limits_{i=1}^x\sum\limits_{j=1}^id_i\\ &=\sum\limits_{i=1}^xd_i\times (x-i+1)\\ &=x\sum\limits_{i=1}^xd_i-\sum\limits_{i=1}^xd_i\times (i-1) \end{aligned} $$ 这样,$r$ 从 $x-\operatorname{lowbit}(x)+1$ 跳到 $x$,$t$ 就要加上: $$ \begin{aligned} \small\sum\limits_{i=1}^xg_i-\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}g_i&\small=\left[x\sum\limits_{i=1}^xd_i-\sum\limits_{i=1}^xd_i\times (i-1)\right]-\left\{[x-\operatorname{lowbit}(x)]\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}d_i-\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}d_i\times (i-1)\right\}\\&=\small\left[x\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}d_i+x\sum\limits_{i=x-\operatorname{lowbit}(x)+1}^xd_i-\sum\limits_{i=1}^xd_i\times (i-1)\right]-\left\{[x-\operatorname{lowbit}(x)]\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}d_i-\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}d_i\times (i-1)\right\}\\&\small=\left\{x\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}d_i-[x-\operatorname{lowbit}(x)]\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}d_i\right\}+x\sum\limits_{i=x-\operatorname{lowbit}(x)+1}^xd_i-\left[\sum\limits_{i=1}^xd_i\times (i-1)-\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}d_i\times (i-1)\right]\\&\small=\operatorname{lowbit}(x)\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}d_i+x\sum\limits_{i=x-\operatorname{lowbit}(x)+1}^xd_i-\sum\limits_{i=x-\operatorname{lowbit}(x)+1}^xd_i\times (i-1) \end{aligned} $$ 注意,$r$ 从 $x-\operatorname{lowbit}(x)+1$ 跳到 $x$, 即 $r$ **跳之前**是 $x-\operatorname{lowbit}(x)+1$,**跳之后**是 $x$。 言归正传,假设刚才说的两个树状数组分别叫 $a,b$, 由上面说的树状数组性质可得,$a_i=\sum\limits_{j=i-\operatorname{lowbit}(x)+1}^i d_i$,$b_i=\sum\limits_{j=i-\operatorname{lowbit}(x)+1}^i d_i\times (i-1)$。 分析一下最后得到 $t$ 在 $r$ 跳之后所加的数值: $$ \operatorname{lowbit}(x)\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}d_i+x\sum\limits_{i=x-\operatorname{lowbit}(x)+1}^xd_i-\sum\limits_{i=x-\operatorname{lowbit}(x)+1}^xd_i\times (i-1) $$ 1. $\operatorname{lowbit}(x)\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}d_i$: 不太好处理,先不管。 2. $x\sum\limits_{i=x-\operatorname{lowbit}(x)+1}^xd_i$:等于 $a_x\times x$。 3. $\sum\limits_{i=x-\operatorname{lowbit}(x)+1}^xd_i\times (i-1)$:等于 $b_x$。 为了处理第 $1$ 项,我们还需要查询 $d_i$ 任意前缀和,但复杂度就不对了。 所以我们可以在 $r$ **每跳一次之后**累加 $a_r$ 到一个变量 $s$ 里, 不难发现,在每次 $r$ **跳之前** $s=\sum\limits_{i=1}^{x-\operatorname{lowbit}(x)}d_i$,正好是我们需要的, 则 $r$ **跳之前**,第一项可以表示为 $s\times \operatorname{lowbit}(x)$,最终结果可以表示为: $$ s\times \operatorname{lowbit}(x)+a_x\times x-b_x $$ 上面强调只能在 $r$ **跳之前**用 $s$ 的值,也就是应该在 $r$ **跳之后**更新 $s$。 同样,直到 $r$ 不能再跳,$r+1$ 就是答案。 ```cpp //c,d 是上面的 a,b int kth(int k) { int r = 0, t = 0, s = 0; for(int i = log值域, x, y, z, cx;i >= 0;--i) { z = 1 << i;if((x = r + z) > 值域) continue; cx = c[x];y = t + s * z + cx * x - d[x]; if(y < k) r = x, t = y, s += cx; } return r + 1; } ``` 参考代码: ```cpp #include <cstdio> #define int long long int n, c[1000050], d[1000050]; void chg(int x, int k) {for(int i = x;i <= 1e6;i += i & -i) c[i] += k, d[i] += (x - 1) * k;} int ask(int x) {int r = 0;for(int i = x;i;i &= i - 1) r += c[i] * x - d[i];return r;} int kth(int k) { int r = 0, t = 0, s = 0; for(int i = 20, x, y, z, cx;i >= 0;--i) { z = 1 << i;if((x = r + z) > 1e6) continue; cx = c[x];y = t + s * z + cx * x - d[x]; if(y < k) r = x, t = y, s += cx; } return r + 1; } signed main() { scanf("%lld", &n); for(int i = 0, o, x, y;i < n;++i) { scanf("%lld%lld", &o, &x); switch(o) { case 1: scanf("%lld", &y);chg(x, 1);chg(y + 1, -1);break; case 2: scanf("%lld", &y);chg(x, -1);chg(y + 1, 1);break; case 3: printf("%lld\n", ask(x - 1) + 1);break; case 4: printf("%lld\n", kth(x));break; case 5: printf("%lld\n", kth(ask(x - 1)));break; case 6: printf("%lld\n", kth(ask(x) + 1));break; } } return 0; } ``` # 关于树套树 没错,这玩意还能套娃。以[树套树板题](https://www.luogu.com.cn/problem/P3380)为例: >维护一个有序数列,5 个操作: >1. 查询 $k$ 在区间内的排名 >2. 查询区间内排名为 $k$ 的值 >3. 修改某一位值上的数值 >4. 查询 $k$ 在区间内的前驱(**若不存在输出 `-2147483647`**) >5. 查询 $k$ 在区间内的后继(**若不存在输出 `2147483647`**) 但是需要解决一些问题。 ## 动态开点 有人就疑惑了:树状数组又不能动态开点,你这空间不是 $O(n^2) $? 咱们用**哈希表**动态开点。~~只能应付随机数据,故意要卡容易 TLE~~ 因为自己写容易锅掉,所以用 pbds。 ```cpp #include <ext/pb_ds/hash_policy.hpp> #include <ext/pb_ds/assoc_container.hpp> using namespace __gnu_pbds; gp_hash_table<int, int> c; //定义一个哈希表,操作和 map 差不多 ``` 实测这玩意比 unordered_map 快了不只两倍。 当然,如果用哈希表,查询和修改就要相应地优化一下: ```cpp struct BIT { __gnu_pbds::gp_hash_table<int, int> c; void chg(int x, int k) {for(;x <= p;x += x & -x) if((c[x] += k) == 0) c.erase(x);} int ask(int x) {int r = 0;for(;x;x &= x - 1) if(c.find(x) != c.end()) r += c[x];return r;} int get(int x) {return c.find(x) != c.end() ? c[x] : 0;} //返回 c[x],之后的 kth 会用到 }c[50050]; ``` ## 整体加减 众所周知,喜闻乐见的 BIT 套动态开点权值线段树是 $O(n\log^2n)$ 的, 少一只 $\log$ 的关键就在于求 $\operatorname{kth}$ 时,线段树的整体加减。 实际上,树状数组也有类似的性质: ```cpp int kth(int l, int r, int k) { int s = 0, t = 0, n = l, m = r; for(int i = lg, x, y;i >= 0;--i) { if((x = s + (1 << i)) > p) continue;--l;y = t; for(;r > l;r &= r - 1) y += c[r].get(x); //c[r].get(x) 即 c[r].c[x] for(;l > r;l &= l - 1) y -= c[l].get(x); //c[l].get(x) 即 c[l].c[x] if(y < k) s = x, t = y;l = n;r = m; } return v[s]; } ``` 当然,这样做就只能整体离散化一次,而不是对于每颗 BIT 分别离散化。 这里以树状数组套权值树状数组为例: ```cpp #include <cmath> #include <cstdio> #include <algorithm> #include <ext/pb_ds/hash_policy.hpp> #include <ext/pb_ds/assoc_container.hpp> #define inf 2147483647 #define Ask(x) ask(l, r, x) #define Kth(x) kth(l, r, x) #define h(x) lower_bound(v, v + p, x) - v + 1 using namespace std; int n, m, p, lg, a[50050], v[100050]; struct Q{int o, l, r, k;}q[50050]; struct BIT { __gnu_pbds::gp_hash_table<int, int> c; void chg(int x, int k) {for(;x <= p;x += x & -x) if((c[x] += k) == 0) c.erase(x);} int ask(int x) {int r = 0;for(;x;x &= x - 1) if(c.find(x) != c.end()) r += c[x];return r;} int get(int x) {return c.find(x) != c.end() ? c[x] : 0;} }c[50050]; void ins(int x, int k) {for(k = h(k);x <= n;x += x & -x) c[x].chg(k, 1);} void del(int x, int k) {for(k = h(k);x <= n;x += x & -x) c[x].chg(k, -1);} int ask(int l, int r, int x) { int s = 0, t = h(x) - 1;--l; for(;r > l;r &= r - 1) s += c[r].ask(t); for(;l > r;l &= l - 1) s -= c[l].ask(t); return s + 1; } int kth(int l, int r, int k) { int s = 0, t = 0, n = l, m = r; for(int i = lg, x, y;i >= 0;--i) { if((x = s + (1 << i)) > p) continue;--l;y = t; for(;r > l;r &= r - 1) y += c[r].get(x); for(;l > r;l &= l - 1) y -= c[l].get(x); if(y < k) s = x, t = y;l = n;r = m; } return v[s]; } int main() { scanf("%d%d", &n, &m); for(int i = 1;i <= n;++i) scanf("%d", a + i), v[p++] = a[i]; for(int i = 0, o, l, r, k;i < m;++i) { scanf("%d%d%d", &o, &l, &r); o == 3 ? v[p++] = r : scanf("%d", &k); q[i] = {o, l, r, k}; } sort(v, v + p);lg = log2(p = unique(v, v + p) - v); for(int i = 1;i <= n;++i) ins(i, a[i]); for(int i = 0, o, l, r, k, t;i < m;++i) { o = q[i].o, l = q[i].l, r = q[i].r, k = q[i].k; switch(o) { case 1: printf("%d\n", Ask(k));break; case 2: printf("%d\n", Kth(k));break; case 3: del(l, a[l]);ins(l, a[l] = r);break; case 4: printf("%d\n", (t = Ask(k)) == 1 ? -inf : Kth(t - 1));break; case 5: printf("%d\n", (t = Ask(k + 1)) > r - l + 1 ? inf : Kth(t));break; } } return 0; } ``` 这里外层树的 `ask` 和 `kth` 的特殊写法,参考[这里](https://www.luogu.com.cn/blog/countercurrent-time/qian-tan-shu-zhuang-shuo-zu-you-hua)的查询优化。