算法背景与问题引入
在日常开发中，我们经常遇到这样的需求：给定一个长度为N的数组，需要多次查询某个子区间内第K小的元素是多少。例如，在日志分析系统中查询某个时间段内第K频繁的操作类型，或在成绩管理系统中查询某班级第K名的分数。
如果数组是静态不变的，我们可以使用可持久化线段树（Persistent Segment Tree，又称主席树）在O(log N)时间内回答每次查询，预处理时间为O(N log N)。这一数据结构的核心思想是：每次单点修改时，不是直接修改原树，而是创建一条从根到叶子的新路径，共享未修改的子树，从而形成一棵”版本树”。
本文将完整讲解主席树的构建原理，并用Java实现静态区间第K大查询和版本回溯功能。
核心原理：函数式持久化
为什么需要可持久化
普通线段树支持单点修改和区间查询，时间复杂度均为O(log N)。但如果我们需要在修改后保留历史版本（例如实现”撤销”操作或查询历史状态的区间第K大），直接复制整棵树的代价是O(N)，无法接受。
主席树的关键洞察在于：一次单点修改只会影响从根节点到目标叶子节点路径上的O(log N)个节点。因此，我们只需新建这O(log N)个节点，其余节点直接复用旧版本的指针，即可在O(log N)时间和空间内创建一个新版本。
版本链的结构
假设我们有一个数组[4, 1, 3, 2]，主席树的构建过程如下：
- 首先对所有数值进行离散化（排序去重），将原始值映射到
[1, M]的整数范围 - 依次将每个元素插入到线段树中，每次插入产生一个新版本
- 版本i的根节点代表前i个元素构成的多重集合
- 查询区间
[L, R]的第K大时，只需用版本R的根减去版本(L-1)的根，即可得到该区间的统计信息
算法步骤详解
第一步：离散化处理
由于数组元素的值域可能很大（如10^9），但元素个数N有限，我们先对所有值进行排序去重，将每个原始值映射到一个紧凑的排名。这使得线段树的范围变为[1, M]，其中M是不重复元素的个数。
第二步：动态开点建线段树
主席树采用动态开点策略：初始时不创建任何节点，每次插入一个值时，从当前版本的根出发，沿途创建需要的新节点。每个节点维护两个信息：
left/right：左右子节点的引用count：当前区间内插入的元素个数
第三步：区间第K大查询
设root[R]为处理完前R个元素后的版本根节点，root[L-1]为前L-1个元素的版本根节点。查询区间[L, R]第K大的过程如下：
- 从两个版本的根节点同时出发
- 计算左子树中元素个数之差：
cnt = left[R].count - left[L-1].count - 如果
K <= cnt，说明第K大的元素在左子树，递归进入左子树 - 否则，第K大的元素在右子树，令
K -= cnt，递归进入右子树 - 当到达叶子节点时，对应的离散化值即为答案
Java完整实现
import java.util.*;
/**
* 可持久化线段树（主席树）实现
* 支持：静态区间第K大查询、版本回溯
*/
public class PersistentSegmentTree {
/**
* 线段树节点类
* 采用动态开点，每个节点记录左右子节点索引和区间内元素个数
*/
static class Node {
int left; // 左子节点在nodes数组中的索引，0表示空
int right; // 右子节点在nodes数组中的索引，0表示空
int count; // 当前区间内的元素个数
Node() {
this.left = 0;
this.right = 0;
this.count = 0;
}
}
// 节点池，动态开点复用数组
private final Node[] nodes;
private int nodeCnt; // 当前已分配的节点数量
private final int maxNodes;
// 每个版本的根节点索引
private final int[] roots;
// 离散化映射：原始值 -> 排名
private final Map<Integer, Integer> compress;
// 反离散化：排名 -> 原始值
private final int[] values;
private final int m; // 离散化后的值域大小
/**
* 构造函数：初始化主席树
* @param arr 原始数组
* @param maxVersions 最大版本数（即数组长度）
*/
public PersistentSegmentTree(int[] arr, int maxVersions) {
int n = arr.length;
// 离散化处理
int[] sorted = arr.clone();
Arrays.sort(sorted);
// 去重得到唯一值列表
int uniqueCount = 0;
for (int i = 0; i < sorted.length; i++) {
if (i == 0 || sorted[i] != sorted[i - 1]) {
uniqueCount++;
}
}
this.values = new int[uniqueCount + 1]; // 1-based
this.compress = new HashMap<>();
int idx = 1;
for (int i = 0; i < sorted.length; i++) {
if (i == 0 || sorted[i] != sorted[i - 1]) {
values[idx] = sorted[i];
compress.put(sorted[i], idx);
idx++;
}
}
this.m = uniqueCount;
// 预分配节点池：每次插入最多新增 log(m) 个节点
// 总节点数上限约为 n * (log(m) + 2)
this.maxNodes = maxVersions * 25;
this.nodes = new Node[maxNodes + 10];
for (int i = 0; i < nodes.length; i++) {
nodes[i] = new Node();
}
this.nodeCnt = 0;
// 初始化版本根数组
this.roots = new int[maxVersions + 2];
roots[0] = 0; // 版本0表示空树
// 依次构建每个版本
for (int i = 1; i <= n; i++) {
int rank = compress.get(arr[i - 1]);
roots[i] = update(roots[i - 1], 1, m, rank);
}
}
/**
* 单点更新：在旧版本基础上插入一个排名为rank的值，返回新版本根节点
* @param prevRoot 旧版本根节点索引
* @param l 当前节点代表的区间左端点
* @param r 当前节点代表的区间右端点
* @param rank 插入值的离散化排名
* @return 新版本根节点索引
*/
private int update(int prevRoot, int l, int r, int rank) {
// 创建新节点，复制旧节点信息
int cur = ++nodeCnt;
nodes[cur].left = nodes[prevRoot].left;
nodes[cur].right = nodes[prevRoot].right;
nodes[cur].count = nodes[prevRoot].count + 1;
if (l < r) {
int mid = (l + r) >>> 1;
if (rank <= mid) {
// 进入左子树，递归更新左子节点
nodes[cur].left = update(nodes[prevRoot].left, l, mid, rank);
} else {
// 进入右子树，递归更新右子节点
nodes[cur].right = update(nodes[prevRoot].right, mid + 1, r, rank);
}
}
return cur;
}
/**
* 查询区间 [L, R] 中第 K 小的元素（1-based，K=1表示最小）
* @param L 区间左端点（1-based，对应原数组下标）
* @param R 区间右端点（1-based）
* @param K 查询第K小
* @return 第K小的原始值
*/
public int queryKth(int L, int R, int K) {
if (K < 1 || K > R - L + 1) {
throw new IllegalArgumentException("K超出范围");
}
return values[query(roots[L - 1], roots[R], 1, m, K)];
}
/**
* 递归查询第K小的离散化排名
* @param leftRoot 前L-1个元素版本的根
* @param rightRoot 前R个元素版本的根
* @param l 当前区间左端点
* @param r 当前区间右端点
* @param K 查询第K小
* @return 离散化排名
*/
private int query(int leftRoot, int rightRoot, int l, int r, int K) {
if (l == r) {
return l;
}
int mid = (l + r) >>> 1;
// 计算左子树中区间[L,R]的元素个数
int leftCount = nodes[nodes[rightRoot].left].count - nodes[nodes[leftRoot].left].count;
if (K <= leftCount) {
// 第K小在左子树
return query(nodes[leftRoot].left, nodes[rightRoot].left, l, mid, K);
} else {
// 第K小在右子树
return query(nodes[leftRoot].right, nodes[rightRoot].right, mid + 1, r, K - leftCount);
}
}
/**
* 获取指定版本的根节点（用于版本回溯）
* @param version 版本号（0表示空集，i表示前i个元素）
* @return 该版本根节点索引
*/
public int getVersionRoot(int version) {
if (version < 0 || version >= roots.length) {
throw new IllegalArgumentException("版本号超出范围");
}
return roots[version];
}
/**
* 基于已有版本创建新版本（支持版本回溯后的继续操作）
* @param baseVersion 基础版本号
* @param value 要插入的原始值
* @return 新版本号
*/
public int forkVersion(int baseVersion, int value) {
if (!compress.containsKey(value)) {
throw new IllegalArgumentException("值不在原始数组中，无法插入");
}
// 注：实际应用中需要动态扩展roots数组
int rank = compress.get(value);
int newRoot = update(roots[baseVersion], 1, m, rank);
// 此处简化处理，实际应维护一个可扩展的版本列表
return newRoot;
}
/**
* 获取某个版本中元素的总个数
* @param version 版本号
* @return 元素个数
*/
public int getVersionSize(int version) {
return nodes[roots[version]].count;
}
// ==================== 主程序与测试 ====================
public static void main(String[] args) {
// 测试数据
int[] arr = {4, 1, 3, 2, 5, 3, 1, 4, 2, 5};
int n = arr.length;
System.out.println("原始数组: " + Arrays.toString(arr));
System.out.println();
// 构建主席树
PersistentSegmentTree pst = new PersistentSegmentTree(arr, n);
// 测试1：查询整个数组的第K小
System.out.println("=== 测试1：查询整个数组 [1, 10] 的第K小 ===");
for (int k = 1; k <= n; k++) {
int result = pst.queryKth(1, n, k);
System.out.printf("第 %d 小的元素是: %d%n", k, result);
}
System.out.println();
// 测试2：查询子区间的第K小
System.out.println("=== 测试2：查询子区间 [3, 7] 的第K小 ===");
// 子数组为 [3, 2, 5, 3, 1]，排序后为 [1, 2, 3, 3, 5]
int[] subArray = Arrays.copyOfRange(arr, 2, 7);
Arrays.sort(subArray);
System.out.println("子数组 [3, 7] 排序后: " + Arrays.toString(subArray));
for (int k = 1; k <= 5; k++) {
int result = pst.queryKth(3, 7, k);
System.out.printf("区间 [3, 7] 第 %d 小的元素是: %d%n", k, result);
}
System.out.println();
// 测试3：查询单元素区间
System.out.println("=== 测试3：查询单元素区间 [5, 5] ===");
int result = pst.queryKth(5, 5, 1);
System.out.printf("区间 [5, 5] 第 1 小的元素是: %d (期望: %d)%n", result, arr[4]);
System.out.println();
// 测试4：版本信息
System.out.println("=== 测试4：各版本的元素个数 ===");
for (int i = 0; i <= n; i++) {
System.out.printf("版本 %d (前 %d 个元素): 元素个数 = %d%n",
i, i, pst.getVersionSize(i));
}
System.out.println();
// 测试5：区间第1小和第K大
System.out.println("=== 测试5：区间 [2, 8] 的最小值、最大值、中位数 ===");
int len = 8 - 2 + 1; // 7个元素
int minVal = pst.queryKth(2, 8, 1);
int maxVal = pst.queryKth(2, 8, len);
int median = pst.queryKth(2, 8, (len + 1) / 2);
System.out.printf("区间 [2, 8] 最小值: %d%n", minVal);
System.out.printf("区间 [2, 8] 最大值: %d%n", maxVal);
System.out.printf("区间 [2, 8] 中位数: %d%n", median);
}
}
复杂度分析
| 操作 | 时间复杂度 | 空间复杂度 |
|---|---|---|
| 初始化（构建所有版本） | O(N log M) | O(N log M) |
| 区间第K大查询 | O(log M) | O(1) |
| 单点更新（创建新版本） | O(log M) | O(log M) |
其中N为数组长度，M为离散化后的值域大小（M ≤ N）。
应用场景拓展
-
历史数据查询：在数据库或日志系统中，每个版本对应一个时间点，可以查询任意历史时刻的数据分布情况。
-
区间中位数查询：通过查询第
(len+1)/2小的元素，可在O(log N)时间内获取区间中位数。 -
带修改的区间第K大：结合树状数组或线段树套主席树，可以支持单点修改后的区间第K大查询（动态主席树）。
-
版本控制系统：在文本编辑器或代码版本管理中，利用主席树的版本链特性，可以快速回溯到任意历史状态。
总结
主席树通过函数式持久化的思想，巧妙地利用动态开点和节点复用，在仅增加O(log N)空间开销的前提下保存了完整的历史版本信息。对于静态区间第K大这类经典问题，它提供了最优的在线查询方案。理解主席树的关键在于把握”路径共享”和”版本相减”这两个核心思想，这也是许多高级数据结构的设计精髓所在。
