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

权值树状数组是啥？

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

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

（如果值域大，可以把操作离线下来然后离散化）

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

怎么搞？

以平衡树板题为例：

维护一个集合，6 个操作：

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

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

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

插入 x

$x$ 的出现次数 $+1$，即修改 $g_x\leftarrow g_x+1$。

删除 x

$x$ 的出现次数 $-1$，即修改 $g_x\leftarrow g_x-1$。

求 x 的排名

不难得出，$\left(\sum _{i=1}^{x-1}{g}_i\right)+1$ 即为所求。

求排名为 x 的数

设 $x$ 的排名表示为 $f(x)=\left(\sum _{i=1}^{x-1}{g}_i\right)+1$。

形式化地，求 $\underset{f(k)\le x}{max}k$，我们把下面的条件移个项：$$\begin{array}{r}\left(\sum _{i=1}^{k-1}{g}_i\right)+1\le x\\ \sum _{i=1}^{k-1}{g}_i\le x-1\end{array}$$

看起来像二分，不过复杂度变成 ${\log }^{2}n$，不太行。

有一个叫树状数组上倍增的 trick（可以去看这道题的题解），可以解决这个问题。

设当前位置为 $r$（初值为 $0$），$t=\sum _{i=1}^r{g}_i$。

不断尝试把 $r$ 往后跳 $2^{\log n},2^{\log n-1},2^{\log n-2},...,2^0$，时刻保证 $\sum _{i=1}^r{g}_i\le x-1$，即 $t<x$。

考虑如何更新 $t$。不难发现，$r$ 只会从 $r-\mathrm{lowbit}(r)+1$ 跳到 $r$。

于是我们直接利用树状数组的结构 $c_r=\sum _{i=r-\mathrm{lowbit}(r)+1}^r{g}_i$，$t\leftarrow t+c_r$ 即可。

$r$ 不能再往后跳，就说明 $r$ 达到满足条件的最大值，

对照上面的公式，$r=k-1$，答案即为 $r+1$。

举个例子：集合内有 $8$ 个数 $1,2,3,4,5,6,7,8$：

现在要求排名为 $6$ 的数。

尝试把 $r$ 向后跳 $2^3$ 到 $r=8$：

$t=8>6$，不能跳，$r=0$，$t=0$。

尝试把 $r$ 向后跳 $2^2$ 到 $r=4$：

$t=4<6$，可以跳，$r=4$，$t=4$。

尝试把 $r$ 向后跳 $2^1$ 到 $r=6$：

$t=6=6$，不能跳，$r=4$，$t=4$。

尝试把 $r$ 向后跳 $2^0$ 到 $r=5$：

$t=5<6$，可以跳，$r=5$，$t=5$。

答案即为 $r+1=6$。

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 的前驱

前驱就是排名为 $\underset{f(i)<f(x)}{max}f(i)$ 的数。

要保证 $f(i)<f(x)$ 的情况下 $f(i)$ 最大，可以令 $f(i)=f(x)-1$。

那么只需要求 $\mathrm{kth}(f(x)-1)$ 即可（注意括号）。

求 x 的后继

同理，后继就是排名为 $\underset{i>x}{min}f(i)$ 的数。

要保证 $i>x$ 的情况下 $f(i)$ 最小，可以令 $i=x+1$。

即求 $\mathrm{kth}(f(x+1))$ （注意括号）。

---

参考代码：

#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$ 的排名为 $\underset{i=1}{\overset{x-1}{max}}c(i)$，其中 $c(i)$ 为 $i$ 在集合中出现的次数。

在此排名意义下，维护平衡树的 6 个操作。允许 ${\log }^{2}n$ 插入删除。值域 $[1,{10}^6]$。

做法类似，不展开叙述：

#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 问题，

但维护 MEX 的树状数组并不是（狭义的）权值树状数组，本文不做说明。

扩展篇

考虑维护一种数据结构：

1. 插入 $i\in [x,y]$。
2. 删除 $i\in [x,y]$。
3. 求 $x$ 的排名。
4. 求排名为 $x$ 的数。
5. 求 $x$ 的前驱。
6. 求 $x$ 的后继。

权值线段树可以轻松解决这种问题，但众所周知 BIT 常数更优。

（注意到这里二者都可以 $\log n$ 维护所有操作）

我们知道，树状数组维护区间修区间查用两个数组，

一个维护差分数组 $d_i$，一个维护 $d_i\times (i-1)$。

那么查询前缀和就可以这样：$$\begin{array}{rl}\sum _{i=1}^x{g}_i& =\sum _{i=1}^x\sum _{j=1}^i{d}_i\\ & =\sum _{i=1}^x{d}_i\times (x-i+1)\\ & =x\sum _{i=1}^x{d}_i-\sum _{i=1}^x{d}_i\times (i-1)\end{array}$$ 这样，$r$ 从 $x-\mathrm{lowbit}(x)+1$ 跳到 $x$，$t$ 就要加上：$$\begin{array}{rl}\sum _{i=1}^x{g}_i-\sum _{i=1}^{x-\mathrm{lowbit}(x)}{g}_i& =[x\sum _{i=1}^x{d}_i-\sum _{i=1}^x{d}_i\times (i-1)]-\{[x-\mathrm{lowbit}(x)]\sum _{i=1}^{x-\mathrm{lowbit}(x)}{d}_i-\sum _{i=1}^{x-\mathrm{lowbit}(x)}{d}_i\times (i-1)\}\\ & =[x\sum _{i=1}^{x-\mathrm{lowbit}(x)}{d}_i+x\sum _{i=x-\mathrm{lowbit}(x)+1}^x{d}_i-\sum _{i=1}^x{d}_i\times (i-1)]-\{[x-\mathrm{lowbit}(x)]\sum _{i=1}^{x-\mathrm{lowbit}(x)}{d}_i-\sum _{i=1}^{x-\mathrm{lowbit}(x)}{d}_i\times (i-1)\}\\ & =\{x\sum _{i=1}^{x-\mathrm{lowbit}(x)}{d}_i-[x-\mathrm{lowbit}(x)]\sum _{i=1}^{x-\mathrm{lowbit}(x)}{d}_i\}+x\sum _{i=x-\mathrm{lowbit}(x)+1}^x{d}_i-[\sum _{i=1}^x{d}_i\times (i-1)-\sum _{i=1}^{x-\mathrm{lowbit}(x)}{d}_i\times (i-1)]\\ & =\mathrm{lowbit}(x)\sum _{i=1}^{x-\mathrm{lowbit}(x)}{d}_i+x\sum _{i=x-\mathrm{lowbit}(x)+1}^x{d}_i-\sum _{i=x-\mathrm{lowbit}(x)+1}^x{d}_i\times (i-1)\end{array}$$ 注意，$r$ 从 $x-\mathrm{lowbit}(x)+1$ 跳到 $x$，

即 $r$ 跳之前是 $x-\mathrm{lowbit}(x)+1$，跳之后是 $x$。

言归正传，假设刚才说的两个树状数组分别叫 $a,b$，

由上面说的树状数组性质可得，$a_i=\sum _{j=i-\mathrm{lowbit}(x)+1}^i{d}_i$，$b_i=\sum _{j=i-\mathrm{lowbit}(x)+1}^i{d}_i\times (i-1)$。

分析一下最后得到 $t$ 在 $r$ 跳之后所加的数值：$$\mathrm{lowbit}(x)\sum _{i=1}^{x-\mathrm{lowbit}(x)}{d}_i+x\sum _{i=x-\mathrm{lowbit}(x)+1}^x{d}_i-\sum _{i=x-\mathrm{lowbit}(x)+1}^x{d}_i\times (i-1)$$

1. $\mathrm{lowbit}(x)\sum _{i=1}^{x-\mathrm{lowbit}(x)}{d}_i$： 不太好处理，先不管。
2. $x\sum _{i=x-\mathrm{lowbit}(x)+1}^x{d}_i$：等于 $a_x\times x$。
3. $\sum _{i=x-\mathrm{lowbit}(x)+1}^x{d}_i\times (i-1)$：等于 $b_x$。

为了处理第 $1$ 项，我们还需要查询 $d_i$ 任意前缀和，但复杂度就不对了。

所以我们可以在 $r$ 每跳一次之后累加 $a_r$ 到一个变量 $s$ 里，

不难发现，在每次 $r$ 跳之前 $s=\sum _{i=1}^{x-\mathrm{lowbit}(x)}{d}_i$，正好是我们需要的，

则 $r$ 跳之前，第一项可以表示为 $s\times \mathrm{lowbit}(x)$，最终结果可以表示为：$$s\times \mathrm{lowbit}(x)+a_x\times x-b_x$$ 上面强调只能在 $r$ 跳之前用 $s$ 的值，也就是应该在 $r$ 跳之后更新 $s$。

同样，直到 $r$ 不能再跳，$r+1$ 就是答案。

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

参考代码：

#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;
}

关于树套树

没错，这玩意还能套娃。以树套树板题为例：

维护一个有序数列，5 个操作：

1. 查询 $k$ 在区间内的排名
2. 查询区间内排名为 $k$ 的值
3. 修改某一位值上的数值
4. 查询 $k$ 在区间内的前驱（若不存在输出 -2147483647）
5. 查询 $k$ 在区间内的后继（若不存在输出 2147483647）

但是需要解决一些问题。

动态开点

有人就疑惑了：树状数组又不能动态开点，你这空间不是 $O(n^2)$？

咱们用哈希表动态开点。只能应付随机数据，故意要卡容易 TLE

因为自己写容易锅掉，所以用 pbds。

#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 快了不只两倍。

当然，如果用哈希表，查询和修改就要相应地优化一下：

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 }^{2}n)$ 的，

少一只 $\log$ 的关键就在于求 $\mathrm{kth}$ 时，线段树的整体加减。

实际上，树状数组也有类似的性质：

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 分别离散化。

这里以树状数组套权值树状数组为例：

#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 的特殊写法，参考这里的查询优化。