从主席树到小波矩阵

· · 算法·理论

摘要

区间第 k 小、区间内小于 x 的数的个数、区间内某个值出现次数，这几类问题在国内 OI 训练里一般会被归到主席树、可持久化线段树或者离线树状数组里。这样做当然没问题，而且主席树确实很经典。

不过如果经常看 Codeforces、ICPC 选手的代码库，或者翻一些国外资料，会遇到另一个名字：小波树，或者更常写成数组形式的小波矩阵。它处理的是同一类静态区间值域统计问题，能在 O(n\,B) 的预处理后，用 O(B) 回答第 k 小、区间排名、值出现次数等询问。其中 B 是值域二进制位数。

小波矩阵本质上就是：把序列按二进制位一层一层稳定划分，同时记录每层前缀里有多少个 1。理解这一点之后，它其实比想象中好写很多。

引入

考虑一个静态数组 a_1,a_2,\dots,a_n，有很多询问，每次给出 l,r，可能要问：

如果只看第一个问题，很容易想到主席树。每个前缀建一棵权值线段树，查询 [l,r] 时用第 r 个版本减去第 l-1 个版本。往左子树还是右子树走，由左子树里的数有多少个决定。

这个思路很顺，但代码里要维护 rt、ls、rs、sum，还要离散化，节点数组也要估好大小。平时写熟了还好，一旦题目本身有别的细节，主席树的代码量就有点压人。

小波矩阵换了一个角度。它不显式建很多棵权值线段树，而是直接把原序列按二进制位分层重排。每一层只做一件事：当前位是 0 的放前面，当前位是 1 的放后面，并且保持同一类内部的相对顺序不变。

这个稳定划分，就是整个结构的核心。

从主席树的判断说起

主席树求区间第 k 小时，每到一个值域节点，就会问一句：区间里有多少个数落在左半边？

如果这个数量不少于 k，答案在左半边；否则答案在右半边。

小波矩阵做的判断很像。只不过这里的“左半边”和“右半边”变成了当前二进制位的 0 和 1。

从最高位往最低位看。当前位为 0 的数一定排在当前位为 1 的数前面，所以如果区间内当前位为 0 的数足够多，第 k 小就会走 0 分支；否则走 1 分支，并把这些 0 分支的数跳过去。

和主席树相比，这里没有节点，也没有版本。每一层只需要知道：某个前缀里有多少个 1。

设第 d 层的前缀和数组是 s_d，那么区间 [l,r] 中当前位为 1 的数量是：

one=s_{d,r}-s_{d,l-1}

当前位为 0 的数量是：

zero=(r-l+1)-one

这一层划分后，所有 0 在前，所有 1 在后。设这一层里 0 的总数是 mid_d。那么区间往下一层映射时：

走 0 分支，新区间为：

[l-s_{d,l-1},\ r-s_{d,r}]

走 1 分支，新区间为：

[mid_d+s_{d,l-1}+1,\ mid_d+s_{d,r}]

刚开始看这两个式子会有点不舒服，建议按含义记。

去 0 分支时，要删掉前面那些被分到 1 的元素。去 1 分支时，要先跳过这一层所有 0，再加上前缀中已经出现过的 1。

只要每层都是稳定划分，原来区间里的元素到了下一层后仍然是一段连续区间，所以查询可以一直往下做。

建立小波矩阵

假设所有数满足：

0\le a_i<2^B

从第 B-1 位到第 0 位依次处理。每一层的流程是：

这里有一个容易忽略的点：一定要稳定划分。不能随便排序，也不能只按这一位粗暴打乱。因为后面的查询依赖“区间映射后仍然连续”这个性质。

查询第 k 小

现在考虑 kth(l,r,k)，其中 k 从 1 开始。

在第 d 层，先算出 zero。如果：

k\le zero

说明答案当前位是 0，把区间映射到 0 分支。

否则答案当前位是 1，令：

k\gets k-zero

同时给答案加上：

2^d

然后进入 1 分支。

这样从高位走到低位，答案的每一位都被确定了。这个过程和主席树求第 k 小几乎是同一个逻辑，只是计数方式不同。

查询小于 x 的个数

再看 rank_less(l,r,x)，也就是查询区间 [l,r] 中小于 x 的数有多少个。

仍然从高位到低位考虑。

如果 x 当前位是 0，那么答案里当前位也只能是 0。因为一旦这一位是 1，就已经大于等于 x 了。所以此时只进入 0 分支，不增加答案。

如果 x 当前位是 1，那么当前位为 0 的所有数都已经严格小于 x，可以直接把 zero 加到答案里。当前位为 1 的数还要继续看后面的位，于是进入 1 分支。

最后得到：

\lvert\{i\mid l\le i\le r,\ a_i<x\}\rvert

于是区间内值落在 [L,R] 的个数也很好写：

\operatorname{rank\_less}(l,r,R+1)-\operatorname{rank\_less}(l,r,L)

值 x 的出现次数就是上式取 L=R=x。

和主席树的对比

两者复杂度其实很接近。

数据结构 预处理 单次查询 空间 更顺手的地方
主席树 O(nB) O(B) O(nB) 国内题常见，扩展到版本思路自然
小波矩阵 O(nB) O(B) O(nB) 接口统一，不用动态开点

小波矩阵比较舒服的地方是，很多操作都能落到 kth 和 rank_less 上。它不像主席树那样需要一直想着“两个版本相减”，也不用担心节点开少了。

当然，它也有明显限制。它最自然的形态是静态序列。如果题目有单点修改、区间修改，那就不适合直接套这个结构了。还有就是它的下标映射一开始很容易写错，第一次写建议拿小数组手推一遍。

所以它不是用来替代主席树的。更准确地说，它是另一个处理静态范围秩问题的模板。主席树会写，小波矩阵也会写，遇到题时选择就会多一点。

代码

下面给一份普通 OI 风格的实现。默认 0\le a_i<2^{30}。如果有负数，可以先离散化；如果值域更大，把 LOG 改大。

#include<bits/stdc++.h>
using namespace std;

const int N=200005;
const int LOG=30;
const int LIM=1<<30;

int n,q;
int a[N],cur[N],tmp[N];
int sum[LOG][N],mid[LOG];

inline void build(){
    for(int i=1;i<=n;i++) cur[i]=a[i];
    for(int d=LOG-1;d>=0;d--){
        int c0=0,c1=0;
        sum[d][0]=0;
        for(int i=1;i<=n;i++){
            int bit=(cur[i]>>d)&1;
            sum[d][i]=sum[d][i-1]+bit;
            if(!bit) tmp[++c0]=cur[i];
        }
        mid[d]=c0,c1=c0;
        for(int i=1;i<=n;i++) if((cur[i]>>d)&1) tmp[++c1]=cur[i];
        for(int i=1;i<=n;i++) cur[i]=tmp[i];
    }
}

inline int kth(int l,int r,int k){
    int ret=0;
    for(int d=LOG-1;d>=0;d--){
        int one=sum[d][r]-sum[d][l-1];
        int zero=r-l+1-one;
        if(k<=zero){
            l=l-sum[d][l-1];
            r=r-sum[d][r];
        }else{
            ret|=1<<d,k-=zero;
            l=mid[d]+sum[d][l-1]+1;
            r=mid[d]+sum[d][r];
        }
    }
    return ret;
}

inline int rank_less(int l,int r,int x){
    if(x<=0) return 0;
    if(x>=LIM) return r-l+1;
    int ret=0;
    for(int d=LOG-1;d>=0;d--){
        int one=sum[d][r]-sum[d][l-1];
        int zero=r-l+1-one;
        if((x>>d)&1){
            ret+=zero;
            l=mid[d]+sum[d][l-1]+1;
            r=mid[d]+sum[d][r];
        }else{
            l=l-sum[d][l-1];
            r=r-sum[d][r];
        }
        if(l>r) break;
    }
    return ret;
}

inline int count_x(int l,int r,int x){
    return rank_less(l,r,x+1)-rank_less(l,r,x);
}

int main(){
    ios::sync_with_stdio(0);
    cin.tie(0);
    cin>>n>>q;
    for(int i=1;i<=n;i++) cin>>a[i];
    build();
    while(q--){
        int op,l,r,x;
        cin>>op>>l>>r>>x;
        if(op==1) cout<<kth(l,r,x)<<'\n';
        else if(op==2) cout<<rank_less(l,r,x)<<'\n';
        else cout<<count_x(l,r,x)<<'\n';
    }
    return 0;
}

正确性说明

建结构时，每一层都按照当前二进制位做稳定划分。因此，一个区间里所有走向同一分支的元素，在下一层仍然是一段连续区间。查询时维护的 [l,r] 始终表示“原询问区间中还没被低位区分的那些数，在当前层的位置”。

求第 k 小时，当前位为 0 的数一定排在当前位为 1 的数前面。若 zero\ge k，答案在 0 分支；否则答案在 1 分支，并跳过这 zero 个更小的数。逐位确定后得到的就是区间第 k 小。

求小于 x 的个数时，如果 x 当前位是 0，只能继续走 0 分支；如果 x 当前位是 1，当前位为 0 的数全部可以计入答案，当前位为 1 的数继续比较后面的位。这个过程按最高不同位分类，刚好统计了所有小于 x 的数。

count_x 由小于 x+1 的个数减去小于 x 的个数得到，所以也是正确的。

复杂度分析

设值域二进制位数为 B。

建结构时每一层扫一遍数组，时间复杂度为：

O(nB)

每次查询只走 B 层，时间复杂度为：

O(B)

空间主要是每层的前缀和数组，复杂度为：

O(nB)

如果取 B=30，它和常见主席树的复杂度在同一档。

适用场景

小波矩阵适合序列静态、询问很多、询问内容和区间值域排名有关的题。比较常见的关键词有：

如果题目一看就是动态修改，那别硬套。小波矩阵最强的地方是静态查询。把它放进工具箱，是为了遇到范围秩问题时多一种写法，而不是见到主席树就全部换掉。

参考资料