详细揭秘：如何发明小波矩阵

· · 算法·理论

以静态区间第 k 小为例。

首先假装你会归并树。归并树是啥？其实就是对归并排序的过程建树。先建立一个线段树，再自底向上归并，求出每个节点对应的区间 [l,r] 的所有元素构成的有序序列。

归并树可以慢速二维数点。具体地，我们要查区间 [l,r] 有多少个 \le x 的数，可以直接找到每个线段树节点，对节点上的序列二分一下。时间复杂度是 \mathcal{O}(\log^2 n) 的。要是直接拿这玩意二分求区间第 k 小，可以做到高贵的 \mathcal{O}(\log^2 n\log V) 复杂度，这也太菜了。

究其原因是没有利用好归并树能做二维数点。不妨换维，将下标与值域互换，相当于求：值域在 [l,r] 内的从左到右第 k 个数。这就成功把要二分的东西放到了线段树维护的维度上，那么可以把二分换成线段树二分做到 \mathcal{O}(\log n\log V)。

还能再给力一点吗？注意到换维过后值域是 \mathcal{O}(n) 的了，如果能 \mathcal{O}(n) 地对同一层处理出来每个数前面有几个数，那么就能再去掉一个 \log n。于是我们开始思考怎么干这个事情。

方便起见把线段树换成 \text{01trie}，这样可以省掉一些讨论。现在每一层的元素个数是 \mathcal{O}(n)，值域也是 \mathcal{O}(n)，我们当然希望能把同一层的所有数都搓到一起。于是不妨将所有左子树扔到最前面，右子树扔到最后面，然后把序列拼起来。大概是这样：

这个大概叫“蝴蝶变换”。方便起见把左子树方向称为红色，右子树方向称为蓝色。

这个树的结构就可以扔掉了，因为我只需要知道一个位置前面有几个蓝色元素，就能直接推出它在蝴蝶变换后会到哪个地方。现在在线段树二分上的过程可以写成这样：

我们甚至只需要知道元素个数！所以每一层做一个前缀和即可。代码大概长这样：

inline 
int kth(int l, int r, int k) {
    int res = 0; l--;
    for (int i = 29; ~i; i--) {
        int ql = b[i][l], qr = b[i][r], c = (r - qr) - (l - ql);
        if (c >= k) l -= ql, r -= qr;
        else res |= 1 << i, k -= c, l = p[i] + ql, r = p[i] + qr;
    }
    return res;
}

其中 b_i 是第 i 层的前缀和，p_i 是第 i 层的红色元素个数。这样我们就做到了 \mathcal{O}(\log V) 查询。

还能再再再给力一点吗？时间复杂度上很难优化了，但是空间 \mathcal{O}(n\log V) 还是不够看。发现我们只要算 1 的个数，可以压个位，空间复杂度做到 \mathcal{O}\left(\frac{n\log V}{w}\right)=\mathcal{O}(n)。

然后你就吊打主席树了。恭喜你发明了小波矩阵！

struct Bitset {
    int n; vector<pair<ull, int>> b;
    Bitset() {}
    Bitset(vector<ull> a) {
        n = (a.size() >> 6) + 1, b.resize(n);
        for (int i = 0; i < a.size(); i++) b[i >> 6].first |= a[i] << (i & 63);
        for (int i = 1; i < n; i++) b[i].second = b[i - 1].second + __builtin_popcountll(b[i - 1].first);
    }
    inline int ask(int k) { return b[k >> 6].second + __builtin_popcountll(b[k >> 6].first & (1ull << (k & 63)) - 1); }
};

struct Wavelet {
    int n; array<Bitset, 30> b; array<int, 30> p;
    Wavelet(vector<int> a) {
        n = a.size(); vector<int> t(n);
        for (int i = 29; ~i; i--) {
            vector<ull> tmp(n); int x = 0, y = 0;
            for (int j = 0; j < n; j++) tmp[j] = a[j] >> i & 1;
            for (int j = 0; j < n; j++) (a[j] >> i & 1 ? t[y++] : a[x++]) = a[j];
            b[i] = Bitset(tmp), p[i] = x;
            for (int j = 0; j < y; j++) a[x + j] = t[j];
        }
    }
    inline 
    int kth(int l, int r, int k) {
        int res = 0; l--;
        for (int i = 29; ~i; i--) {
            int ql = b[i].ask(l), qr = b[i].ask(r), c = (r - qr) - (l - ql);
            if (c >= k) l -= ql, r -= qr;
            else res |= 1 << i, k -= c, l = p[i] + ql, r = p[i] + qr;
        }
        return res;
    }
};

去掉两个 \log n 的思路实际上都还蛮实用的。首先是第一个换维。换维的思想十分实用，如果说遇到要维护两维的信息，其中一维的维护难度比另一维更大，那么可以尝试交换两维。第二个就是将同一层拉平统一处理。这在分散层叠中也有用到，大致思想就是，先有一个复杂度跟总元素个数挂钩的预处理，再通过将所有元素捆一起预处理，从而让复杂度均摊。

事实上以这个结构理解，可以使得对小波矩阵进行的拓展更加直观。比如说你可以看出来，小波矩阵本来就是用来做二维数点的，求区间第 k 大只是换维一下顺手的事，所以改一改就能过掉 P4148 了（