Wavelet Matrix

lgswdn_SA · 2024-06-24 16:03:39 · 题解

Wavelet Matrix

考虑这样的一个结构（以下序列全部都是 0-index）：

对于原序列 a，设其最高位是 H。令 c_{H+1,i}=a_i。然后我们从高到低遍历每一位 x。
对于 x，我们先将 c_{x} 继承自 c_{x+1}，然后我们记录 c_x 每个元素二进制最高位 x 的值 b_{x,i}。
然后我们对原来的 c_x 做一次变换得到新的 c_{x}：把第 x 位为 0 的元素挪到左边，然后第 x 位为 1 的元素挪到右边。设 p_x 表示这一位 0 的个数。
eg：a=[1,3,2,0]，我们先得到 b_1=[0,1,1,0]，然后得到 c_1=[1,0,3,2]，p_1=2；然后我们得到 b_0=[1,0,1,0]，然后得到 c_0=[0,2,1,3]，p_0=2。注意 c 并不是一个有序结构。

- 从高往低遍历每一位 $x$。 - 考虑 $i\in[l,r]$ 中 $b_{x,i}=0$ 的个数 $w$。如果 $w<k$ 那么意味着这一位必须是 $0$，那么下一层我们只需要在满足 $b_{x,i}=0$ 的前缀里面问询即可。具体而言，设 $l'$ 表示这一层下标 $<l$ 的 $b_{x,j}=0$ 的个数，$r'$ 表示这一层下标 $\le r$ 的 $b_{x,j}=0$ 的个数，那么只需要递归到下一层时将 $[l,r]$ 替换成 $[l',r'-1]$ 即可。 - 如果 $w>k$ 那么意味着这一位必须是 $1$，那么我们贡献给答案，再将 $k$ 减掉 $w$ 后递归进满足 $b_{x,i}=1$ 的后缀里面问询即可。具体而言，设 $l'$ 表示这一层下标 $<l$ 的 $b_{x,j}=1$ 的个数，$r'$ 表示这一层下标 $\le r$ 的 $b_{x,j}=1$ 的个数，那么只需要递归进到下一层时将 $[l,r]$ 替换成 $[l'+p_x,r'+p_x-1]$ 即可。注意到我们并不需要对于每个 $x$ 都维护出 $c$，只需在构建时动态维护变换即可。而 $b$ 的维护可以使用一个手写的 bitset（支持单点修改，以及 $O(1)$ 区间查 popcount）。空间复杂度 $O(\frac{n\log V}{w})$。实现部分参考自 https://zhuanlan.zhihu.com/p/590974585 （如在构建 Wavelet Matrix 时使用 `std::stable_partition` 等，同时该文章内有引用自 https://miti-7.hatenablog.com/entry/2018/04/28/152259 的一些视觉化的演示）。 ``` const int N=(1<<18)+5,H=29; int n,a[N],m; struct bst { ull b[N/64+5],c[N/64+5]; int m; #define msk(i) ((1ull<<((i)&63))-1) void init(int n) {m=n/64+1;} void set(int i) {b[i>>6]|=1ull<<(i&63);} bool get(int i) {return (b[i>>6]>>(i&63))&1;} int p1(int i) {return c[i>>6]+popcnt(b[i>>6]&msk(i&63));} int p0(int i) {return i-c[i>>6]-popcnt(b[i>>6]&msk(i&63));} int s1(int l,int r) {return p1(r+1)-p1(l);} int s0(int l,int r) {return p0(r+1)-p0(l);} void getcnt() {rep(i,1,m-1) c[i]=c[i-1]+popcnt(b[i-1]);} }; namespace WT { bst b[30]; int p[30],c[N]; void build() { rep(i,0,n-1) c[i]=a[i]; per(h,29,0) { b[h].init(n); rep(i,0,n-1) if(c[i]&(1<<h)) b[h].set(i); b[h].getcnt(); p[h]=stable_partition(c,c+n,[&](int x) { return !((x>>h)&1);})-c; } } int qry(int l,int r,int k,int x=0) { k--; per(h,29,0) { int cnt=b[h].s0(l,r); if(cnt>k) l=b[h].p0(l), r=b[h].p0(r+1)-1, assert(r<=p[h]); else x|=1<<h, l=p[h]+b[h].p1(l), r=p[h]+b[h].p1(r+1)-1, k-=cnt; } return x; } } signed main() { n=read(), m=read(); rep(i,0,n-1) a[i]=read(); WT::build(); while(m--) { int l=read()-1, r=read()-1, k=read(); printf("%d\n",WT::qry(l,r,k)); } return 0; } ```