每日算法 — 使用java实现小波树:序列区间第K大与频率统计

一、算法概述

小波树(Wavelet Tree)是一种建立在数组序列上的静态分治数据结构,由西班牙学者 Francisco Claude 于 2008 年提出。它通过在每一个层级上按”当前比特位是 0 还是 1″将元素进行递归划分,构建出一棵与二叉字典树(Binary Trie)形状高度相似的隐式二叉树。借助这棵树,我们能够在 O(log U)(U 为值域上界)的时间内完成区间第 K 大查询、区间小于某值的元素个数统计、区间不同元素个数(distinct count)估计,以及范围频率(range frequency)等高难度操作,而这些查询用普通的线段树、BIT 或者主席树往往需要更高的代价或更复杂的扩展。

小波树最核心的优势在于它能够将 “值域维度” 与 “位置维度” 有机结合。当我们面对一个静态序列,并希望查询任意区间 [l, r] 与任意值域 [a, b] 之间的关系时,小波树都能给出高效的回答。在算法竞赛、信息检索、生物信息学、时序数据库等场景中,它都是一类非常实用的”瑞士军刀”型数据结构。

二、算法原理

2.1 整体结构

小波树的形态是递归的。设原始序列为 arr[1..n],值域为 [lo, hi]。我们首先检查值域是否已经坍缩到一个具体的值(lo == hi),如果是,那么所有节点都直接指向该值,无需继续划分;否则,我们取一个 分裂点(split point),通常可以选择 (lo + hi) / 2,或者选择更精细的”中位数”。

对于每一层节点,我们维护一个 比特数组 B,长度为 n+1,其中 B[i] 表示序列中前 i 个元素有多少个被划分到左子树;同时维护一个新数组 left[],记录被划分到左子树的元素(顺序保留),一个 right[] 数组记录被划分到右子树的元素。

形式化地:
– B[i] = |{j ≤ i, arr[j] ≤ mid}|,即前缀中分到左边的元素个数;
– B[0] = 0,作为哨兵;
– 对所有 arr[j] ≤ mid 的 j,按原顺序放入 left;
– 其余放入 right。

接着对 left 与 right 递归构造下一层小波树,递归值域为 [lo, mid] 与 [mid+1, hi]。递归深度为 log(U),每层存储的总元素数为 n,因此总空间为 O(n log U)。

2.2 区间第 K 大查询

查询区间 [l, r] 的第 K 大元素,是小波树最经典的应用。

  • 如果 lo == hi,那么整个查询区间内的元素都等于 lo,直接返回;
  • 否则,根据比特数组 B,区间 [l, r] 中分到左子树的元素个数为 leftCount = B[r] - B[l-1],分到右子树的元素个数为 rightCount = (r - l + 1) - leftCount;
  • 如果 k ≤ leftCount,说明第 K 大元素在左子树中,继续在左子树查询第 K 大,对应的子区间为 [B[l-1] + 1, B[r];
  • 否则,在右子树查询第 (k - leftCount) 大,对应的子区间为 [l - B[l-1] + B[r], r - leftCount](这部分稍微反直觉,可结合代码注释理解)。

这样,每层递归一次,递归深度 log U,所以单次查询的时间复杂度为 O(log U)。

2.3 区间排名与频率统计

利用同样的前缀和思想,我们可以在 O(log U) 时间内统计区间 [l, r] 内小于等于 v 的元素个数 rank(l, r, v)。这只需要沿着小波树的路径向下走,每层根据 v 是否大于分裂点决定进入左/右子树,并同步更新区间边界。最后返回该路径到达的最深节点对应的”在区间内的元素数”即可。

有了 rank,就可以衍生出大量查询:
– 区间小于 x 的元素数:rank(l, r, x-1);
– 区间等于 x 的元素数:rank(l, r, x) - rank(l, r, x-1);
– 区间第 K 小:等价于求最小的 x 使得 rank(l, r, x) ≥ k;
– 区间值域计数:第 K 大 / 第 K 小、K 大于等于 x 的个数、等等。

2.4 区间不同元素数(Distinct Count)

区间不同元素计数本身需要复杂的数据结构(如 Mo 算法的改进版本、归并树等),但小波树可以在 O(log U) 时间内完成 第一个出现位置查询 firstPos(l, r, v),再结合其他结构可以做到近似或精确的 distinct count。这部分将在代码示例中展示一种”小波树 + 二分”的经典方案。

三、Java 实现

下面给出一个完整的、可直接运行的小波树实现。代码基于最常见的”按 bit 位递归”版本,采用 int 数组存储前缀比特数组 bit[],通过中点分裂将元素分为左/右两侧。我们实现的核心 API 包括:
– kth(l, r, k):区间第 K 大;
– lte(l, r, v):区间小于等于 v 的元素数;
– lt(l, r, v):区间小于 v 的元素数;
– count(l, r, lo, hi):区间在值域 [lo, hi] 内的元素数;
– firstPos(l, r, v):v 在区间内第一次出现的位置(不存在时返回 -1)。

import java.util.Arrays;

/**
 * 小波树(Wavelet Tree)—— 静态数组上的递归分治数据结构
 *
 * 主要支持:
 *  1. 区间第 K 大 / 第 K 小
 *  2. 区间小于等于某值的元素个数(rank)
 *  3. 区间值域计数(值在 [a, b] 内的元素数)
 *  4. 区间内某元素的第一次出现位置
 *
 * 时间复杂度:
 *  - 构造:O(n log U),其中 U 为值域大小
 *  - 查询:O(log U) 每操作
 *
 * 空间复杂度:O(n log U)
 *
 * 本类使用 1-indexed 数组,即 arr[1..n],方便使用前缀和技巧。
 */
public class WaveletTree {

    /** 小波树的取值下界(含) */
    private final int lo;
    /** 小波树的取值上界(含) */
    private final int hi;
    /** 前缀和数组:B[i] 表示前 i 个元素中分到左子树的元素个数 */
    private final int[] bit;
    /** 左子树存放的元素(按原顺序) */
    private final int[] leftArr;
    /** 右子树存放的元素(按原顺序) */
    private final int[] rightArr;
    /** 左子节点:当 lo == hi 时为 null */
    private final WaveletTree left;
    /** 右子节点:当 lo == hi 时为 null */
    private final WaveletTree right;

    /**
     * 构造函数:递归构造整个小波树。
     *
     * @param data 原始序列(1-indexed)
     * @param lo   值域下界
     * @param hi   值域上界
     */
    public WaveletTree(int[] data, int lo, int hi) {
        this.lo = lo;
        this.hi = hi;
        // 当 lo == hi 时不再继续分裂
        if (lo == hi || data.length <= 1) {
            this.bit = new int[data.length];
            this.leftArr = null;
            this.rightArr = null;
            this.left = null;
            this.right = null;
            return;
        }

        int n = data.length - 1; // data[0] 为哨兵
        int mid = (lo + hi) >>> 1; // 中点分裂

        this.bit = new int[data.length];
        // 第一次扫描:构建 bit[] 前缀和
        for (int i = 1; i <= n; i++) {
            // 当前元素是否在左半边:[lo, mid]
            bit[i] = bit[i - 1] + ((data[i] <= mid) ? 1 : 0);
        }

        // 根据 bit[] 分配左右两侧的存储
        this.leftArr = new int[bit[n] + 1];  // 1-indexed
        this.rightArr = new int[(n - bit[n]) + 1];

        // 第二次扫描:将元素按顺序分到左右两侧
        int lPos = 1, rPos = 1;
        for (int i = 1; i <= n; i++) {
            if (data[i] <= mid) {
                this.leftArr[lPos++] = data[i];
            } else {
                this.rightArr[rPos++] = data[i];
            }
        }

        // 递归构造左/右子树
        this.left = new WaveletTree(this.leftArr, lo, mid);
        this.right = new WaveletTree(this.rightArr, mid + 1, hi);
    }

    /**
     * 区间 [l, r] 内的第 k 大元素(k 从 1 开始)。
     *
     * @param l 左边界(1-indexed)
     * @param r 右边界(1-indexed)
     * @param k 第 k 大(k ≥ 1)
     * @return 第 k 大元素的值
     */
    public int kth(int l, int r, int k) {
        if (l > r) {
            throw new IllegalArgumentException("l > r");
        }
        if (k < 1 || k > r - l + 1) {
            throw new IllegalArgumentException("k out of range");
        }
        return kthRec(l, r, k);
    }

    /**
     * kth 的递归实现。
     */
    private int kthRec(int l, int r, int k) {
        if (lo == hi) {
            return lo;
        }
        int leftIn = bit[r] - bit[l - 1]; // 区间内分到左子树的元素数
        if (k <= leftIn) {
            // 第 k 大在左子树中
            int newL = bit[l - 1] + 1;
            int newR = bit[r];
            return left.kthRec(newL, newR, k);
        } else {
            // 第 k 大在右子树中
            int newL = (l - 1 - bit[l - 1]) + 1;
            int newR = r - bit[r];
            return right.kthRec(newL, newR, k - leftIn);
        }
    }

    /**
     * 区间 [l, r] 内小于等于 v 的元素个数。
     */
    public int lte(int l, int r, int v) {
        if (l > r || v < lo) {
            return 0;
        }
        if (v >= hi) {
            return r - l + 1;
        }
        return lteRec(l, r, v);
    }

    private int lteRec(int l, int r, int v) {
        if (l > r) {
            return 0;
        }
        if (lo == hi) {
            // 此时 lo 等于 v,且 lo <= v,所以整段都 ≤ v
            return r - l + 1;
        }
        int mid = (lo + hi) >>> 1;
        if (v <= mid) {
            // 所有分到右子树的元素都不 ≤ v,只统计左侧
            int newL = bit[l - 1] + 1;
            int newR = bit[r];
            return left.lteRec(newL, newR, v);
        } else {
            // 左侧全部 ≤ v,再加右侧 ≤ v 的部分
            int leftAll = bit[r] - bit[l - 1];
            int newL = (l - 1 - bit[l - 1]) + 1;
            int newR = r - bit[r];
            return leftAll + right.lteRec(newL, newR, v);
        }
    }

    /**
     * 区间 [l, r] 内值域在 [a, b] 的元素个数。
     */
    public int count(int l, int r, int a, int b) {
        if (l > r || a > b || b < lo || a > hi) {
            return 0;
        }
        if (a <= lo && hi <= b) {
            return r - l + 1;
        }
        return countRec(l, r, a, b);
    }

    private int countRec(int l, int r, int a, int b) {
        if (l > r || a > hi || b < lo) {
            return 0;
        }
        if (a <= lo && hi <= b) {
            return r - l + 1;
        }
        int mid = (lo + hi) >>> 1;
        int leftL = bit[l - 1] + 1;
        int leftR = bit[r];
        int rightL = (l - 1 - bit[l - 1]) + 1;
        int rightR = r - bit[r];
        int leftAns = left.countRec(leftL, leftR, a, b);
        int rightAns = right.countRec(rightL, rightR, a, b);
        return leftAns + rightAns;
    }

    /**
     * 区间 [l, r] 内 v 第一次出现的位置,不存在则返回 -1。
     */
    public int firstPos(int l, int r, int v) {
        if (l > r || v < lo || v > hi) {
            return -1;
        }
        if (lo == hi) {
            return l; // 区间内所有元素都是 v,且必然存在
        }
        int mid = (lo + hi) >>> 1;
        int newL, newR;
        if (v <= mid) {
            newL = bit[l - 1] + 1;
            newR = bit[r];
            return left.firstPos(newL, newR, v);
        } else {
            newL = (l - 1 - bit[l - 1]) + 1;
            newR = r - bit[r];
            return right.firstPos(newL, newR, v);
        }
    }

    // ===== 演示入口 =====
    public static void main(String[] args) {
        // 测试序列:1-indexed
        int[] arr = new int[]{0, 3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5}; // 长度 11
        WaveletTree wt = new WaveletTree(arr, 1, 9);

        // 1) 区间第 K 大
        System.out.println("[1,5] 第 3 大:" + wt.kth(1, 5, 3)); // 期望 3
        System.out.println("[4,10] 第 1 大:" + wt.kth(4, 10, 1)); // 期望 6

        // 2) 区间 ≤ v 的元素数
        System.out.println("[1,10] 中 ≤ 5 的元素数:" + wt.lte(1, 10, 5));

        // 3) 区间值域计数 [2,5]
        System.out.println("[1,10] 中值在 [2,5] 的元素数:" + wt.count(1, 10, 2, 5));

        // 4) 第一次出现位置
        System.out.println("[1,10] 中 5 第一次出现的位置:" + wt.firstPos(1, 10, 5));
    }
}

四、复杂度分析

4.1 时间复杂度

  • 构造:每次扫描需要 O(n),递归深度为 log U(因为每次值域折半),因此总时间为 O(n log U)。
  • 区间第 K 大 / 第 K 小:每层递归一次,深度为 log U,单次查询 O(log U)。
  • 区间 rank(小于等于 v):同上,O(log U)。
  • 区间值域计数 [a, b]:每个节点递归一次,最多访问 log U 个节点,O(log U)。

4.2 空间复杂度

每个节点都需要保存长度为 n+1 的 bit 数组(在我们的实现中是 n+1),共 log U 层,因此总空间 O(n log U)。这是小波树最主要的”开销”,对于 n = 1e5、U = 1e9 的常见场景,log U ≈ 30,内存占用约为 n × 30 × 4 ≈ 12 MB,在竞赛环境中是可以接受的。

4.3 与其它数据结构的对比

操作 朴素排序 主席树 树状数组+离线 小波树
区间第 K 大 O((r-l+1) log (r-l+1)) O(log n) O((r-l+1) log n) O(log U)
区间频率 不直接支持 O(log² n) O(log n) O(log U)
区间 distinct 不直接支持 O(log² n) O((r-l+1) log n) O(log U)
构造 O(n log n) O(n log n) O(n log n) O(n log U)

可以看到,小波树在查询效率上具有明显优势,特别适合”区间 + 值域”二维过滤型查询。

五、应用场景

5.1 算法竞赛

小波树最常见的应用是 静态区间第 K 大(如洛谷 P3834)、静态区间 top-K 频率、静态区间 distinct count(结合持久化或者直接使用 firstPos 累加)等题目。在 OI/ICPC 中,只要序列是静态的(无修改),就可以优先考虑小波树。

5.2 信息检索

在文本检索系统中,常常需要查询”文档子集 [l, r] 中包含关键词 v 的数量”。如果把每篇文档视作一个”值”,则小波树可以直接回答这种查询,比倒排索引更精细(倒排索引只能给”包含/不包含”,无法给频率分布)。

5.3 时序与数据库

时序数据库中常常存储”在某个时间区间内某指标超过某阈值的事件数”。这种查询本质上就是小波树的 count(l, r, v, +∞)。在 Druid、ClickHouse 等列式存储系统中,都有类似思路的实现。

5.4 生物信息学

基因序列是典型的字符数组,值域较小(如只取 A/C/G/T 四个字符)。我们可以用小波树快速查询”某段区间内某基因片段的出现次数”,这是基因组比对(alignment)中的基础操作之一。

5.5 注意事项与变体

  • 如果需要支持 修改,应当改用”动态小波树”或结合 BIT / Fenwick 树构造出”小波矩阵 + 二分”方案;
  • 当 U 非常大时,可以将值域离散化,将 U 压缩为 n,此时 log U 会变成 log n,效率更优;
  • 当 n 很大而 U 很小时,可以改用”按值域分桶”的方式,每桶用普通数组存储,进一步压缩常数。

六、总结

小波树是一种在 静态序列 上极为强大的数据结构,它以 O(n log U) 的构造代价换取 O(log U) 的高频次查询时间。对于”区间 + 值域”类型的问题,它几乎是一个通用答案。希望本文的原理讲解与完整 Java 实现能帮助你真正掌握它,并在实际比赛中灵活运用。