浅谈 Wavelet Matrix

__staring__ · 2025-06-18 19:43:50 · 算法·理论

前几天看到了这个数据结构，觉得比较优雅，记录一下。

由于 wavelet matrix 貌似没有正式的译名，这里将其称为小波矩阵。

总的来说，小波矩阵是一种基于值域二进制分解的静态简洁数据结构，其基于小波树 wavelet tree。

关于简洁数据结构 Succinct data structure，简洁位向量 Succinct Indexable Dictionaries，以及其支持的 \opt{access,rank,select} 操作，由于竞赛中较少使用，这里仅稍作介绍，感兴趣的读者可自行查询相关资料。

在竞赛中，小波矩阵能以 O(n\log{V})-O(\frac{n\log{V}}{\omega}) 的时空复杂度预处理，单次询问 O(\log{V}) 时间，解决静态无权区间 k 小区间排名等问题。

一般认为（word-RAM model）计算机字长 \omega \ge \log{V}，所以在这种语境下完全可以认为小波矩阵做到了 O(n) 的空间复杂度。

约定

下文叙述和代码实现上，对序列和值域位均使用 0-index。

默认序列长度为 n，值域为 [0,2^A),V=2^A。

小波树是一棵 O(\log{V}) 层的二叉树，记根节点为 rt，记 u_{0/1} 表示结点 u 的左/右子节点，h(u) 表示结点 u 的高度。

节点 u 维护 0/1 序列 b_u，b_{u,i} 为某个序列 a'_u 中 a'_{u,i} 第 h(u) 位的值。

具体的递归定义如下：将序列 a'_u 拆分，

若子节点高度 =-1 或其维护序列 =\varnothing 则返回。初始 h(rt)=A-1,a'_{rt}=a，这里的 a 指需要对其构建小波树的序列。

注意小波树并不存储 a'_u，只存储 b_u 的信息。时空复杂度 O(n\log{V})-O(n\log{V})。

不难发现小波树的形态与 01-trie，或者说动态开点线段树非常相似，因此其也支持相关操作。例如查询序列 a 中 <x 的元素个数：从根节点开始，若 x 第 h(u) 位为 1，则将 u_0 的元素加入答案，并进入 u_1，否则直接进入 u_0。

\opt{id}_c(u,k)

若 b_{u,k}=c，求对应的元素在 a'_{u_c} 中的下标。

不难发现就是 b_u[0,k-1] 中 c 的个数，即 \opt{rank}_c(b_u,k-1)。

\opt{access}(a,k)

求 a_k 的值。

从根节点开始，维护 a_k 在 a'_u 的下标 k'，可以知道 a_k 的第 h(u) 位为 b_{u,k'}，其在子节点的下标是 \opt{id}_{b_{u,k'}}(u,k')。

一路向下查找即可得到 a_k 的每一位。当然，也可以在叶节点处额外维护值，直接返回。

\opt{rank}_v(a,k)

求 a[0,k] 中 v 的个数。

同样从根节点开始，记 v 第 h(u) 位为 c。可能存在的 v 一定在 u_c 中，且定义时的子序列保证了相对顺序不变，那么 \opt{rank}_v(a'_u,k)=\opt{rank}_v(a'_{u_c},\opt{id}_c(u,k))。

递归到叶节点时答案为 k，即 \opt{rank}_v(a_\text{leaf},k)=k

\opt{select}_v(a,k)

求 a 中第 k 个 v 的下标。

考虑 v 的第 h(u) 位 c，a'_u 中第 k 个 v 在 u_c 的下标为 k'=\opt{select}_v(a'_{u_c},k)，那么 a'_{u_c} 的第 k' 个元素在 a'_u 的下标为 \opt{select}_c(b_u,k')。

于是 \opt{select}_v(a'_u,k)=\opt{select}_c(b_u,\opt{select}_v(a'_{u_c},k))。对于叶节点，有 \opt{select}_v(a'_\text{leaf},k)=k。

若 \opt{select}_c(b_u,k') 为 O(1) 实现，复杂度为 O(\log{V})。

小波矩阵

小波矩阵的主要思想是，将小波树每一层的所有位向量，按照一定顺序组成一个长为 n 的位向量一起存储，无需维护树形结构。

\opt{build}

记当前要构建序列 a[0,n-1] 的小波矩阵 b[0,A][0,n-1]。为了便于说明，记辅助数组 a'[0,A][0,n-1]。

初始有 b_{A,i}=0,a'_A=a。枚举二进制位 u : A-1 \rarr 0，记录 b_{u,i} 为 a'_{u+1,i} 第 u 位的值，然后将 a'_{u+1} 按照 (b_{u,i},i) 二元组排序得到 a'_u。

可以发现这等价于，对于小波树每一层的位向量，所有 b_{u_0} 一定在所有 b_{u_1} 前面，同时保证 b_{u_0},b_{u_1} 内部在上一层的相对顺序不变。

注意，小波树的叶层节点是有序的，但小波矩阵最终的 a'_0 不一定有序。根据小波树容易发现，在 a'_0 中所有相同的元素一定只形成 1 个区间。

朴素实现的时空复杂度为 O(n\log{V})-O(n\log{V})，这里给出一份便于理解。a' 在实现上是不需要的，只是为了方便下文说明。

void build(int n, int a[])
{
    for(int u = A - 1; ~u; u --)
    {
        for(int i = 1; i <= n; i ++)
            b[u][i] = a[i] >> u & 1;
        stable_partition(a + 1, a + n + 1,
            [&u](int x){return ~x >> u & 1;});
    }
}

\opt{id}_c(a'_u,k)

若 a'_{u,k} 的第 u 位为 c，求其在 a'_{u-1} 的下标。

• 若 c=0，那么只有 a'_u[0,k-1] 中第 u 位为 0 的元素在 a'_{u,k} 之前，即 \opt{id}_0(a'_u,k)=\opt{rank}_0(b_u,k-1)。

• 若 c=1，那么 a'_u 中第 u 位为 0 的所有元素，和 a'_u[0,k-1] 中第 u 位为 1 的元素都在 a'_{u,k} 之前，即 \opt{id}_1(a'_u,k)=\opt{rank}_0(b_u,n-1)+\opt{rank}_1(b_u,k-1)。

为了方便描述 \opt{select}，除了向下转移的 \opt{id}，我们再定义：

\opt{id}^{-1}_c(a'_u,k)

若 a'_{u-1,k} 的第 u 位为 c，求其在 a'_u 的下标（可以认为是 \opt{id} 的逆运算）。

• 若 c=0，那么其为 b_u 的第 k+1 个 0，即 \opt{id}^{-1}_0(a'_u,k)=\opt{select}_0(b_u,k+1)。

• 若 c=1，那么其为 b_u 的第 k-\opt{rank}_0(b_u,n-1)+1 个 1，即 \opt{id}^{-1}_1(a'_u,k)=\opt{select}_1(b_u,k-\opt{rank}_0(b_u,n-1)+1)。

确定如何转移之后，查询的思路与小波树类似。

\opt{access}(a,k)

求 a_k 的值。

初始 u=A-1,k'=k，表示 a_k 在 a'_u 的下标 k'。查询得到 c=\opt{access}(b_u,k')，即 a_k 第 u 位的值，那么转移到 k' \larr \opt{id}_c(a'_u,k'),u \larr u-1。

\opt{rank}_v(a,l,r)

求 a[l,r] 中 v 的个数。

若沿用小波树的思路查询 \opt{rank}_v(a,0,k)，在 a'_0 无法知道对应区间左端点，也就不能确定答案。当然可以选择在 a'_0 维护每个元素段的起点，但我们可以直接维护左端点，解决区间查询。

若 v 的第 u 位为 c，根据小波树可以知道所有的 v 一定在同一个子节点内，由于相对顺序不变，因此可以仅排出 a'_u[l,r] 内第 u 位为 \lnot c 的元素。

初始 u=A-1，逐次更新 l \larr \opt{id}_c(a'_u,l-1)+1, r \larr \opt{id}_c(a'_u,r), u \larr u-1，最终答案即为 r-l+1。

\opt{select}_v(a,l,k)

求 a[l,n-1] 中第 k 个 v 的下标。

同样地，查询 \opt{select}_v(a,0,k) 时若不维护起点，则在 a'_0 无法给出第 k 个元素的下标，也就无法向上计算答案。因此维护小波矩阵上的起点，解决后缀查询。

若 v 的 第 u 位为 c，类似小波树地考虑第 k 个 v 在下一层的下标为 k'=\opt{select}_v(a'_{u-1},\opt{id}_c(a'_u,l-1)+1,k)，那么下一层下标为 k' 的元素在这层的下标为 \opt{id}^{-1}_c(a'_u,k')。

初始 u=A-1，先向下更新 l \larr \opt{id}_c{a'_u,l-1}+1, u \larr u-1，u=0 后令 k \larr l+k-1，再向上更新 k \larr \opt{id}_c^{-1}(a'_u,k), u \larr u+1。

若 \opt{id}^{-1} 中的 \opt{select} 为 O(1) 实现，则复杂度为 O(\log{V})。

优化

在上述基础操作中，维护 b 的内层数据结构需要支持 \opt{access}(b_u,k), \opt{rank}_{0/1}(b_u,k), \opt{select}_{0/1}(b_u,k)。

或许有读者发现，这正好是简洁位向量 Succinct Indexable Dictionaries 支持的操作。然而竞赛中 \opt{access} 非常容易实现，\opt{select} 绝大部分时候用不到，且简洁位向量虽然理论更优，其常数和编码复杂度并不如下面的方法。

我们直接使用 O(\log{V}) 个长 O(\frac{n}{\omega}) 的 bitset 维护 b 即可，或者说将 b 每一行按 \omega=64 进行分块压位存储。

对于 \opt{access}(b_u,k)，简单位运算即可。

对于 \opt{rank}_{0/1}(b_u,k)，块外维护 b_u 每一块的前缀和 s_u，块内使用 __builtin_popcountll，单次复杂度 O(1)。

对于 \opt{select}_{0/1}(b_u,k)，块外块内二分，单次复杂度 O(\log{\frac{n}{\omega}}+\log{\omega})=O(\log{n})。

于是，我们得到了预处理时间复杂度 O(n\log{V})，空间复杂度 O(\frac{n\log{V}}{\omega})/O(n) 的小波矩阵。关于上述基础操作，这里给出笔者的示例代码，供读者参考理解。

struct bitSet
{
    using ULL = unsigned long long;
    #define ppcll __builtin_popcountll

    int len, zero;
    vector <pair <ULL, int>> bit;

    bitSet(int n): len(n >> 6), zero(n)
    {
        bit.resize(len + 1, make_pair(0, 0));
    }

    void set(int k)
    {
        --zero;
        bit[k >> 6].first |= 1ull << (k & 63);
    }

    void init()
    {
        for(int i = 1; i <= len; i ++)
            bit[i].second = bit[i - 1].second + ppcll(bit[i - 1].first);
    }

    int access(int k)
    {
        return bit[k >> 6].first >> (k & 63) & 1;
    }

    int rank1(int k)
    {
        return bit[k >> 6].second + ppcll(bit[k >> 6].first << (~k & 63));
    }

    int rank0(int k)
    {
        return k - rank1(k);
    }

    int binarySearch(int v, int k)
    {
        static constexpr ULL m32 = 0xffffffff;
        static constexpr ULL m16 = 0xffff;
        static constexpr ULL m8  = 0xff;
        static constexpr ULL m4  = 0xf;
        static constexpr ULL m2  = 0x3;
        static constexpr ULL m1  = 0x1;

        int c, pos = 0;
        if(c = ppcll(v & m32); c < k) v >>= 32, k -= c, pos |= 32;
        if(c = ppcll(v & m16); c < k) v >>= 16, k -= c, pos |= 16;
        if(c = ppcll(v & m8 ); c < k) v >>=  8, k -= c, pos |=  8;
        if(c = ppcll(v & m4 ); c < k) v >>=  4, k -= c, pos |=  4;
        if(c = ppcll(v & m2 ); c < k) v >>=  2, k -= c, pos |=  2;
        if(c = ppcll(v & m1 ); c < k) v >>=  1, k -= c, pos |=  1;
        return pos;
    }

    int select1(int k)
    {
        int pos = 0;
        int l = 0, r = len, mid;
        while(l <= r)
        {
            mid = (l + r) >> 1;
            if(bit[mid].second < k)
                l = mid + 1, pos = mid;
            else r = mid - 1;
        }
        k -= bit[pos].second;
        return pos << 6 | binarySearch(bit[pos].first, k);
    }

    int select0(int k)
    {
        int pos = 0;
        int l = 0, r = len, mid;
        while(l <= r)
        {
            mid = (l + r) >> 1;
            if((mid << 6) - bit[mid].second < k)
                l = mid + 1, pos = mid;
            else r = mid - 1;
        }
        k -= (pos << 6) - bit[pos].second;
        return pos << 6 | binarySearch(~bit[pos].first, k);
    }

    #undef ppcll
};

\opt{kth-min}(a,l,r,k)

求 a[l,r] 中第 k 小的元素。

由小波树可以知道，u \rarr u - 1 时 a'_u[l,r] 在 a'_{u-1} 只会重组为两个区间，其中一个第 u 位全 0，另一个第 u 位全 1。于是从 u=A-1 开始，维护答案 v=0：

• 若 a'_u[l,r] 内第 u 位为 0 的元素个数 \ge k，即 \opt{rank}_0(b_u,l,r) \ge k，显然v 这一位为 0，更新 l \larr \opt{id}_0(a'_u,l-1)+1, r \larr \opt{id}_0(a'_u,r), u \larr u-1。

• 否则 v 这一位为 1，去掉这一位为 0 元素的贡献，更新 v \larr v+2^u, l \larr \opt{id}_1(a'_u,l-1)+1, r \larr \opt{id}_1(a'_u,r), k \larr k-\opt{rank}_0(b_u,l,r), u \larr u-1。

\opt{ranking}(a,l,r,v)

求 a[l,r] 中 <v 的元素个数。

类似地，从 u=A-1 开始，维护答案 k=0：

• 若 v 第 u 位为 1，将所有这一位为 0 的元素加入统计，更新 k \larr k+ \opt{rank}_0(a'_u,l,r), l \larr \opt{id}_1(a'_u,l-1)+1, r \larr \opt{id}_1(a'_u,r), u \larr u-1。

• 否则直接更新 l \larr \opt{id}_0(a'_u,l-1)+1, r \larr \opt{id}_0(a'_u,r), u \larr u-1。

实现

通过上面的一些操作可以发现，小波矩阵继承了小波树带来的 01-trie 相似结构，且能通过简洁位向量进行区间转移。相当于提取区间的 01-trie 进行查询，因此大部分持久化 01-trie 或者说主席树支持的操作，在小波矩阵上稍作修改即可实现。

位向量的统一存储和简洁位向量的线性空间对缓存较为友好，这使得小波矩阵可以以小常数或小空间替代主席树。当然，若是统计信息带权，空间复杂度只能做到 O(n\log{V})。

这里给出笔者的实例代码，供读者参考理解。一些建议：

• 朴素实现中的 stable_partition() 虽然可以做到 O(n)，但在没有内部临时缓冲区时，是 O(n\log{n}) 的复杂度，因此一般应该替换为手写版本。

• 有时初始化中的 tmp0 可以不使用，直接在原数组上修改，最后原数组即保存小波矩阵叶层位向量的原始值。

• 所有涉及左端点 l 的操作，都可以将左闭换成左开，来省掉对 l 的偏移。

• 内层 bitset 和外层 wavelet matrix 都使用 0-index，但在查询时使用 1-index，因为内层查 [0,k) 前缀和比 [0,k] 方便很多，查 [l,r] 时放到外层偏移掉 l 再换成右开，变为 [l-1,r]。

struct waveletMatrix
{
    int n, m;
    vector <bitSet> bit;

    waveletMatrix(int _n, int _m, int a[]): n(_n), m(_m)
    {
        bit.resize(m, bitSet(n));
        vector <int> tmp0(a, a + n + 1), tmp1(n + 1);
        for(int u = m - 1; ~u; u --)
        {
            int x = 0, y = 0;
            for(int i = 1; i <= n; i ++)
                if(tmp0[i] >> u & 1)
                    tmp1[++y] = tmp0[i], bit[u].set(i);
                else tmp0[++x] = tmp0[i];
            bit[u].init();
            for(int i = 1; i <= y; i ++)
                tmp0[x + i] = tmp1[i];
        }
    }

    int id(int u, int k, int c)
    {
        return !c ? bit[u].rank0(k) : bit[u].zero + bit[u].rank1(k);
    }

    int di(int u, int k, int c)
    {
        return !c ? bit[u].select0(k) : bit[u].select1(k - bit[u].zero);
    }

    int access(int k)
    {
        int res = 0;
        for(int u = m - 1; ~u; u --)
        {
            int c = bit[u].access(k);
            res |= c << u;
            k = id(u, k, c);
        }
        return res;
    }

    int rank(int l, int r, int v)
    {
        for(int u = m - 1; ~u; u --)
        {
            l = id(u, l - 1, v >> u & 1) + 1;
            r = id(u, r, v >> u & 1);
        }
        return r - l + 1;
    }

    int select(int l, int k, int v)
    {
        for(int u = m - 1; ~u; u --)
            l = id(u, l - 1, v >> u & 1) + 1;
        l += k - 1;
        for(int u = 0; u < m; u ++)
            l = di(u, l, v >> u & 1);
        return l;
    }

    int kth_min(int l, int r, int k)
    {
        int res = 0;
        for(int u = m - 1; ~u; u --)
        {
            int c = bit[u].rank0(r) - bit[u].rank0(l - 1);
            l = id(u, l - 1, c < k) + 1;
            r = id(u, r, c < k);
            if(c < k)
            {
                res |= 1 << u;
                k -= c;
            }
        }
        return res;
    }

    int ranking(int l, int r, int v)
    {
        int res = 0;
        for(int u = m - 1; ~u; u --)
        {
            int c = v >> u & 1;
            if(c) res += bit[u].rank0(r) - bit[u].rank0(l - 1);
            l = id(u, l - 1, c) + 1;
            r = id(u, r, c);
        }
        return res;
    }
};

以上给出了一份标准化实现代码，比较繁杂且有很多不必要内容。以下则是一份简略版实现。

struct bits
{
    int n;
    vector <pair <ULL, int>> bit;
    #define ppc __builtin_popcountll
    bits() {}
    bits(vector <ULL> arr)
    {
        n = (arr.size() >> 6) + 1, bit.resize(n);
        for (int i : viota(0, ssize(arr)))
            bit[i >> 6].first |= arr[i] << (i & 63);
        for (int i : viota(1, n))
            bit[i].second = bit[i - 1].second + ppc(bit[i - 1].first);
    }
    int ask(int k) {return bit[k >> 6].second +
        ppc(bit[k >> 6].first & ((1ull << (k & 63)) - 1));}
    #undef ppc
};

struct wave
{
    int n;
    array <bits, 30> bit;
    array <int, 30> pos;
    wave(vector <int> arr)
    {
        n = arr.size();
        for (int u : viota(0, 30) | vreve)
        {
            bit[u] = bits (arr | vtran([&] (auto x)
                {return x >> u & 1ull;}) | ranges::to <vector> ());
            pos[u] = n - ranges::stable_partition
                (arr, [&] (int x) {return ~x >> u & 1;}).size();
        }
    }
    int kthmin(int l, int r, int k)
    {
        int v = 0; --l;
        for (int u : viota(0, 30) | vreve)
        {
            int ll = bit[u].ask(l), rr = bit[u].ask(r);
            if (int c = r - rr - l + ll; c >= k) l -= ll, r -= rr;
            else v |= 1 << u, k -= c, l = pos[u] + ll, r = pos[u] + rr;
        }
        return v;
    }
    int rankin(int l, int r, int v)
    {
        int k = 0; --l;
        for (int u : viota(0, 30) | vreve)
        {
            int ll = bit[u].ask(l), rr = bit[u].ask(r);
            if (int c = r - rr - l + ll; ~v >> u & 1) l -= ll, r -= rr;
            else k += c, l = pos[u] + ll, r = pos[u] + rr;
        }
        return k;
    }
};

动态

小波矩阵是静态的数据结构，其维护的信息一般满足可减性或者说有逆元，需要插入删除时可以考虑维护一组负贡献的小波矩阵，可考虑几种经典实现：

• 定期重构，代价为单次修改均摊 O(\sqrt{n\log{V}})\sim O(\sqrt{n}\log{V})，支持插入删除。

• 二进制分组，代价为单次修改均摊 O(\log{V}\log{n})，支持插入和特殊删除：若删除不可减且为撤销，考虑斜二进制分组，当存在三组相同大小的组时合并其中两组，删除时仍然重构，否则不支持删除。

• 平衡树维护内层位向量，代价为单次修改 O(\log{V}\log{n})，支持插入删除。

例题

区间 kth-min

P3834 【模板】可持久化线段树 2

区间 kth-min 模板题，直接使用上面的模板即可。

在线写法直接对原数组修改，无需记录二分过程，返回最后端点即可。

离线则是考虑对每一位，把所有询问的移动统一处理，把小波矩阵做到 O(\frac{n}{\omega}) 的空间复杂度，提交时为最优解 rk2。

在线2 和 主席树2 指的是，以题目的最大值域 [0,10^9] 为准，在线 和 主席树 则分别是 [0,\max] 和 [\min,\max] 的值域。

在线 248ms 4.21mb

在线2 367ms 4.95mb

离线 206ms 5.78mb

主席树 399ms 42.46mb

主席树2 574ms 76.53mb

区间 ranking

P3810 【模板】三维偏序

详细题解，使用区间 ranking 配合二进制分组即可。

小波矩阵 980ms 5.32mb

cdq 263ms 6.60mb

主席树(内存池垃圾回收) 1.62s 37.40mb

主席树(vec 动态开点) 1.05s 23.36mb

区间 kth-min，ranking

P14513 [NFLSPC #8] 追忆desuwa

详细题解，拆点后先做一遍 kth-min 再做一遍 ranking。

二分带权信息

P12866 [JOI Open 2025] 抽奖 / Lottery

详细题解，这里讲述一下实现。

要解决的问题为：给定长为 n 的数组 a_i，多次询问，求最大的 K 满足 \frac{(r-l+1)K}{2} \le \sum_{i=l}^r \min(a_i,K)，且已发现可以二分 K。

考虑类似二分答案的思想，维护可能的答案区间 [l,r]，已计算的答案 res，一定小于最终答案的元素之和 sum，一定大于最终答案的元素个数 cnt，小波矩阵由于需要维护元素和，空间复杂度 O(n\log{V})。

记当前在第 u 层，若 res+2^u 满足条件，则 [l,r] 转移到 [\opt{id}_1(a'_u,l-1)+1,\opt{id}_1(a'_u,r)]，并更新 res,sum，否则令 [l,r] 转移到 [\opt{id}_0(a'_u,l-1)+1,\opt{id}_0(a'_u,r)]，并更新 cnt。

若不离散化则以 O(\log{V}) 的时间直接得到答案，否则以 O(\log{n}) 的时间得到原序列中的最优答案，此时 sum,cnt 都已确定，计算不等式的最优解即可。

小波矩阵 2.51s 141.86mb 提交时是最优解 rk3，没打过 rk1 rk2 老哥的主席树。

主席树 2.46s 199.91mb 什么叫我的主席树跑到最优解了？

动态区间 kth-min，ranking

P2617 Dynamic Rankings

P3380 【模板】树套树

P4278 带插入区间K小值

详细题解，把内层的 bitset 换成平衡树，看上去能做一车动态问题。

矩形 kth-min

P1527 [国家集训队] 矩阵乘法

详细题解，树套小波矩阵，不彳亍，小波矩阵套小波矩阵，彳亍。

动态矩形 kth-min

P4848 崂山白花蛇草水

详细题解，二进制分组合并常数太大了，根本跑不过带根号做法（悲。打不过就以后加入。

to be continued...

参考资料

小波矩阵——ShwStone

算法学习笔记：Wavelet Tree 求解区间第K小的杀器(一)——OnjoujiToki