每日算法 — 使用java实现主席树：可持久化线段树与区间第K大查询

引言

主席树（Chairman Tree），又称可持久化线段树或函数式线段树，是一种支持查询历史版本的高级数据结构。它由黄嘉泰（hjt）在OI竞赛中推广，因此得名”主席树”。本文用Java实现主席树，解决静态区间第K大查询问题，并讲解其空间优化与可持久化原理。

一、核心思想

普通线段树每个节点维护区间聚合信息。主席树的关键洞察是：每次单点修改只影响O(log n)个节点，因此可以复用未修改的子树，仅新建受影响的节点链，从而保存所有历史版本。

1.1 节点结构

public class PersistentSegTree {

    /**
     * 动态节点: 左右子节点用下标引用，支持动态开点
     */
    static class Node {
        int left, right;   // 左右子节点在pool中的下标，0表示空
        int count;         // 该区间内元素个数

        Node() {
            this.left = this.right = 0;
            this.count = 0;
        }
    }

    // 节点池，预分配避免频繁GC
    private final Node[] pool;
    private int poolPtr = 0;

    // 每个版本的根节点下标
    private final int[] roots;

    // 离散化后的值域
    private final int[] sortedValues;
    private final int n; // 值域大小

    public PersistentSegTree(int maxNodes, int maxVersions, int[] sortedValues) {
        this.pool = new Node[maxNodes];
        for (int i = 0; i < pool.length; i++) pool[i] = new Node();
        this.roots = new int[maxVersions];
        this.sortedValues = sortedValues;
        this.n = sortedValues.length;
    }

    /**
     * 新建节点，可选复制已有节点的内容
     */
    private int newNode(int cloneFrom) {
        poolPtr++;
        if (cloneFrom != 0) {
            pool[poolPtr].left = pool[cloneFrom].left;
            pool[poolPtr].right = pool[cloneFrom].right;
            pool[poolPtr].count = pool[cloneFrom].count;
        } else {
            pool[poolPtr].left = pool[poolPtr].right = 0;
            pool[poolPtr].count = 0;
        }
        return poolPtr;
    }
}

1.2 离散化

主席树建立在值域上，需先将原数组离散化：

public class PersistentSegTree {

    /**
     * 对数组进行离散化，返回排序去重后的值数组
     */
    public static int[] discretize(int[] arr) {
        int[] sorted = arr.clone();
        Arrays.sort(sorted);
        // 去重
        int m = 0;
        for (int i = 0; i < sorted.length; i++) {
            if (i == 0 || sorted[i] != sorted[i-1]) {
                sorted[m++] = sorted[i];
            }
        }
        return Arrays.copyOf(sorted, m);
    }

    /**
     * 查找值在离散化数组中的位置（1-based）
     */
    private int getRank(int value) {
        int pos = Arrays.binarySearch(sortedValues, value);
        return pos + 1; // 转为1-based，方便线段树处理
    }
}

二、可持久化单点更新

public class PersistentSegTree {

    /**
     * 在prevVersion的基础上，在值域位置pos增加1
     * 返回新版本根节点下标
     * 
     * @param prevRoot 上一版本的根节点下标
     * @param l 当前节点代表的值域左边界
     * @param r 当前节点代表的值域右边界
     * @param pos 要更新的值域位置
     * @return 新版本的根节点下标
     */
    public int update(int prevRoot, int l, int r, int pos) {
        int curRoot = newNode(prevRoot); // 复制旧节点
        pool[curRoot].count++; // 当前区间计数+1

        if (l != r) {
            int mid = (l + r) >>> 1;
            if (pos <= mid) {
                // 更新左子树
                int newLeft = update(pool[prevRoot].left, l, mid, pos);
                pool[curRoot].left = newLeft;
            } else {
                // 更新右子树
                int newRight = update(pool[prevRoot].right, mid + 1, r, pos);
                pool[curRoot].right = newRight;
            }
        }
        return curRoot;
    }

    /**
     * 构建初始版本（空树）
     */
    public int build(int l, int r) {
        int root = newNode(0);
        if (l != r) {
            int mid = (l + r) >>> 1;
            pool[root].left = build(l, mid);
            pool[root].right = build(mid + 1, r);
        }
        return root;
    }

    /**
     * 依次插入数组元素，构建所有版本
     * roots[i]表示前i个元素构成的版本
     */
    public void buildVersions(int[] arr) {
        roots[0] = build(1, n); // 空版本
        for (int i = 0; i < arr.length; i++) {
            int pos = getRank(arr[i]);
            roots[i + 1] = update(roots[i], 1, n, pos);
        }
    }
}

三、区间第K大查询

3.1 核心查询算法

利用两个版本的差分，可在O(log n)内查询任意区间的第K大值：

public class PersistentSegTree {

    /**
     * 查询区间[L,R]中的第k大值
     * 原理: version[R] - version[L-1] 即为区间[L,R]的统计信息
     * 
     * @param leftRoot  version[L-1]的根
     * @param rightRoot version[R]的根
     * @param l 当前值域左边界
     * @param r 当前值域右边界
     * @param k 第k大（k从1开始）
     * @return 第k大值在离散化数组中的下标
     */
    public int queryKth(int leftRoot, int rightRoot, int l, int r, int k) {
        if (l == r) {
            return l; // 找到答案
        }

        int mid = (l + r) >>> 1;

        // 左子树的元素个数 = 右版本左子树count - 左版本左子树count
        int leftCount = pool[pool[rightRoot].left].count - pool[pool[leftRoot].left].count;

        if (k <= leftCount) {
            // 第k大在左子树
            return queryKth(pool[leftRoot].left, pool[rightRoot].left, l, mid, k);
        } else {
            // 第k大在右子树
            return queryKth(pool[leftRoot].right, pool[rightRoot].right, mid + 1, r, k - leftCount);
        }
    }

    /**
     * 对外接口: 查询原数组区间[L,R]（1-based）的第k大值
     */
    public int query(int L, int R, int k) {
        int pos = queryKth(roots[L - 1], roots[R], 1, n, k);
        return sortedValues[pos - 1]; // 转回0-based取原值
    }
}

3.2 查询原理可视化

假设值域为[1,8]，数组前R个元素构成version[R]，前L-1个构成version[L-1]：

version[R]:        version[L-1]:
      [1-8]:5           [1-8]:2
     /       \         /       \
  [1-4]:3   [5-8]:2  [1-4]:1   [5-8]:1
  ...                  ...

区间[L,R]的统计 = version[R] - version[L-1]:
  [1-8]: 5-2 = 3 个元素
  [1-4]: 3-1 = 2 个元素
  [5-8]: 2-1 = 1 个元素

通过比较左右子树的元素个数与k的关系，可快速定位第k大。

四、完整实现与测试

import java.util.*;

public class ChairmanTreeDemo {

    public static void main(String[] args) {
        // 测试数据
        int[] arr = {3, 1, 4, 1, 5, 9, 2, 6, 5, 3};
        int m = arr.length;

        // 离散化
        int[] sorted = PersistentSegTree.discretize(arr);
        System.out.println("离散化值域: " + Arrays.toString(sorted));

        // 估算节点数: 每次update新建约log2(n)个节点
        // 总节点数 ≈ m * (log2(n) + 2) + 4*n
        int maxNodes = m * 25 + sorted.length * 4;
        PersistentSegTree pst = new PersistentSegTree(maxNodes, m + 1, sorted);

        // 构建所有版本
        pst.buildVersions(arr);
        System.out.println("构建完成，总节点数: " + pst.getNodeCount());

        // 查询测试
        System.out.println("\n区间第K大查询测试:");

        // 查询arr[1..5] = {3,1,4,1,5} 的第1大 -> 5
        System.out.println("arr[1..5] 第1大 = " + pst.query(1, 5, 1));

        // 查询arr[1..5] = {3,1,4,1,5} 的第3大 -> 3 (排序后:5,4,3,1,1)
        System.out.println("arr[1..5] 第3大 = " + pst.query(1, 5, 3));

        // 查询arr[1..5] = {3,1,4,1,5} 的第5大 -> 1
        System.out.println("arr[1..5] 第5大 = " + pst.query(1, 5, 5));

        // 查询arr[3..8] = {4,1,5,9,2,6} 的第2大 -> 6
        System.out.println("arr[3..8] 第2大 = " + pst.query(3, 8, 2));

        // 验证: 暴力排序对比
        System.out.println("\n暴力验证:");
        verifyBruteForce(arr, pst, 1, 5, 3);
        verifyBruteForce(arr, pst, 3, 8, 2);
        verifyBruteForce(arr, pst, 1, 10, 5);
    }

    private static void verifyBruteForce(int[] arr, PersistentSegTree pst, int L, int R, int k) {
        int[] sub = Arrays.copyOfRange(arr, L - 1, R);
        Arrays.sort(sub);
        int expected = sub[sub.length - k]; // 第k大
        int actual = pst.query(L, R, k);
        System.out.printf("[%d,%d] 第%d大: 期望=%d, 实际=%d, %s%n",
            L, R, k, expected, actual, expected == actual ? "PASS" : "FAIL");
    }
}

五、空间优化：动态开点与压缩

5.1 节点池预分配策略

public class PersistentSegTree {

    /**
     * 更精确的空间估算
     * 每个元素插入产生约 log2(n) 个新节点
     * 空树有 4*n 个节点（满二叉树）
     */
    public static int estimateNodes(int dataSize, int valueRange) {
        int logN = 32 - Integer.numberOfLeadingZeros(valueRange);
        return dataSize * (logN + 2) + valueRange * 4;
    }

    public int getNodeCount() {
        return poolPtr;
    }
}

5.2 静态数组版（竞赛优化）

对于已知数据规模的场景，可完全使用静态数组避免对象开销：

public class StaticChairmanTree {

    // 用数组下标模拟指针
    static final int MAXN = 200005;   // 最大数据量
    static final int MAXM = MAXN * 25; // 最大节点数

    static int[] left = new int[MAXM];
    static int[] right = new int[MAXM];
    static int[] count = new int[MAXM];
    static int[] roots = new int[MAXN];
    static int ptr = 0;

    static int update(int prev, int l, int r, int pos) {
        int cur = ++ptr;
        left[cur] = left[prev];
        right[cur] = right[prev];
        count[cur] = count[prev] + 1;

        if (l != r) {
            int mid = (l + r) >>> 1;
            if (pos <= mid) {
                left[cur] = update(left[prev], l, mid, pos);
            } else {
                right[cur] = update(right[prev], mid + 1, r, pos);
            }
        }
        return cur;
    }

    static int query(int u, int v, int l, int r, int k) {
        if (l == r) return l;
        int mid = (l + r) >>> 1;
        int leftCount = count[left[v]] - count[left[u]];
        if (k <= leftCount) {
            return query(left[u], left[v], l, mid, k);
        } else {
            return query(right[u], right[v], mid + 1, r, k - leftCount);
        }
    }
}

六、扩展应用

6.1 区间小于K的个数

public class PersistentSegTree {

    /**
     * 查询区间[L,R]中小于等于val的元素个数
     */
    public int countLessOrEqual(int L, int R, int val) {
        int pos = upperBound(sortedValues, val); // 最后一个<=val的位置
        if (pos == 0) return 0;
        return countLE(roots[L-1], roots[R], 1, n, 1, pos);
    }

    private int countLE(int leftRoot, int rightRoot, int l, int r, int ql, int qr) {
        if (ql > r || qr < l) return 0;
        if (ql <= l && r <= qr) {
            return pool[rightRoot].count - pool[leftRoot].count;
        }
        int mid = (l + r) >>> 1;
        return countLE(pool[leftRoot].left, pool[rightRoot].left, l, mid, ql, qr)
             + countLE(pool[leftRoot].right, pool[rightRoot].right, mid + 1, r, ql, qr);
    }

    private static int upperBound(int[] arr, int val) {
        int lo = 0, hi = arr.length;
        while (lo < hi) {
            int mid = (lo + hi) >>> 1;
            if (arr[mid] <= val) lo = mid + 1;
            else hi = mid;
        }
        return lo;
    }
}

6.2 带修改的可持久化线段树（树套树）

如果需要在原数组上支持单点修改并查询历史版本，可结合树状数组+主席树：

public class DynamicChairmanTree {

    // 树状数组的每个节点维护一棵主席树
    // 修改时更新O(log n)棵树状数组节点，每棵更新O(log n)
    // 总复杂度: 修改O(log^2 n)，查询O(log^2 n)

    private final PersistentSegTree[] bitTrees;
    private final int size;

    public DynamicChairmanTree(int n, int[] sortedValues) {
        this.size = n;
        this.bitTrees = new PersistentSegTree[n + 1];
        for (int i = 1; i <= n; i++) {
            bitTrees[i] = new PersistentSegTree(n * 25, n + 1, sortedValues);
        }
    }

    // 树状数组lowbit
    private int lowbit(int x) { return x & (-x); }

    /**
     * 在位置idx增加delta（对应离散化后的值位置）
     */
    public void add(int idx, int pos, int delta) {
        for (int i = idx; i <= size; i += lowbit(i)) {
            // 对每棵树状数组节点的主席树进行更新
            // 此处简化，实际需要维护每个BIT节点的根
        }
    }
}

七、复杂度分析

八、适用场景与局限

适用场景:
– 静态数组的区间第K大/小查询
– 需要查询历史版本的数据结构问题
– 离线处理的区间统计问题

局限:
– 不支持原数组的修改（需改用树套树，复杂度升至O(log²n)）
– 空间开销较大（约为普通线段树的log n倍）
– 值域离散化要求数据可比较

总结

本文实现了主席树（可持久化线段树）解决区间第K大查询问题。核心思想是利用修改的局部性，每次仅克隆O(log n)个节点来保存新版本，通过版本差分实现区间查询。读者可以将其扩展到区间中位数、区间众数、带修改的持久化等更复杂的场景中。