子数组中位数查询问题及Java实现优化咨询

给定一个包含N个元素的数组A，对于A的任意长度为len的子数组，将其元素按非降序排序后，处于位置ceil(len/2)（数组及子数组均采用1索引）的元素称为该子数组的中位数。需要编写程序处理Q次查询，每次查询给定两个整数L和R，求子数组A_L、A_{L+1}…A_R的中位数。

输入格式 #

• 第一行：整数N
• 第二行：N个空格分隔的整数（表示数组A）
• 第三行：整数Q
• 接下来Q行：每行两个空格分隔的整数L和R

输出格式 #

对于每个查询，输出对应子数组的中位数。

约束条件 #

• 1 ≤ N, Q ≤ 5×10^4
• 1 ≤ A_i ≤ 10^9
• 1 ≤ L ≤ R ≤ N

示例 #

输入：

6
2 4 5 3 1 6
3
1 5
2 4
3 6

输出：

3
4
5

低效实现问题 #

最初的实现对每个查询的子数组单独排序后找中位数，时间复杂度为O(Q*K log K)（K为查询子数组的长度），当Q和K都达到5e4时，运算量会突破1e10级别，完全无法通过时间限制。核心代码如下：

private static final int generous = 1;
private static int median(int[] arr) {
    Arrays.sort(arr);
    int mid = arr.length / 2;
    if (mid + mid == arr.length) {
        return (arr[mid-1] + arr[mid] + generous) / 2;
    } else {
        return arr[mid];
    }
}

private static int[] getMedian(int[] arr) {
    int[] result = new int[arr.length];
    for (int i = 0; i < arr.length; i++) {
        result[i] = median(Arrays.copyOfRange(arr, 0, i+1));
    }
    return result;
}

需要将时间复杂度优化到O(N log N + Q log N)级别，以下是两种可行的方案：

方案一：归并树 + 二分答案 #

思路 #

1. 离散化数组：将原数组元素去重排序，映射到较小的索引范围，解决数值范围过大的问题。
2. 构建归并树：线段树的每个节点存储对应区间的排序数组，方便快速统计区间内≤某个值的元素数量。
3. 二分中位数：对每个查询，通过二分离散化后的数值范围，结合归并树的统计结果，找到满足条件的最小数值，即为中位数。

代码实现 #

import java.io.*;
import java.util.*;

public class MedianQuery {
    static int[][] tree;
    static int[] arr;
    static int[] sorted;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int N = Integer.parseInt(br.readLine());
        arr = new int[N];
        StringTokenizer st = new StringTokenizer(br.readLine());
        for (int i = 0; i < N; i++) {
            arr[i] = Integer.parseInt(st.nextToken());
        }

        // 离散化处理
        sorted = arr.clone();
        Arrays.sort(sorted);
        int M = 0;
        for (int i = 0; i < N; i++) {
            if (i == 0 || sorted[i] != sorted[i-1]) {
                sorted[M++] = sorted[i];
            }
        }
        sorted = Arrays.copyOf(sorted, M);

        // 构建归并树
        int size = 1;
        while (size < N) size <<= 1;
        tree = new int[2 * size][];
        // 填充叶子节点
        for (int i = 0; i < N; i++) {
            tree[size + i] = new int[]{arr[i]};
        }
        // 构建内部节点
        for (int i = size - 1; i > 0; i--) {
            tree[i] = merge(tree[2 * i], tree[2 * i + 1]);
        }

        // 处理查询
        int Q = Integer.parseInt(br.readLine());
        StringBuilder sb = new StringBuilder();
        while (Q-- > 0) {
            st = new StringTokenizer(br.readLine());
            int L = Integer.parseInt(st.nextToken()) - 1; // 转换为0索引
            int R = Integer.parseInt(st.nextToken()) - 1;
            int len = R - L + 1;
            int k = (len + 1) / 2; // 中位数对应的元素个数要求

            // 二分找中位数
            int left = 0, right = M - 1;
            while (left < right) {
                int mid = (left + right) / 2;
                int cnt = query(L, R, sorted[mid]);
                if (cnt >= k) {
                    right = mid;
                } else {
                    left = mid + 1;
                }
            }
            sb.append(sorted[left]).append('\n');
        }
        System.out.print(sb);
    }

    // 合并两个有序数组
    static int[] merge(int[] a, int[] b) {
        int[] res = new int[a.length + b.length];
        int i = 0, j = 0, idx = 0;
        while (i < a.length && j < b.length) {
            res[idx++] = a[i] <= b[j] ? a[i++] : b[j++];
        }
        while (i < a.length) res[idx++] = a[i++];
        while (j < b.length) res[idx++] = b[j++];
        return res;
    }

    // 查询[L,R]中<=x的元素个数
    static int query(int L, int R, int x) {
        int res = 0;
        L += tree.length / 2;
        R += tree.length / 2;
        while (L <= R) {
            if (L % 2 == 1) {
                res += upperBound(tree[L], x);
                L++;
            }
            if (R % 2 == 0) {
                res += upperBound(tree[R], x);
                R--;
            }
            L >>= 1;
            R >>= 1;
        }
        return res;
    }

    // 二分查找数组中<=x的元素个数
    static int upperBound(int[] arr, int x) {
        int left = 0, right = arr.length;
        while (left < right) {
            int mid = (left + right) / 2;
            if (arr[mid] > x) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }
        return left;
    }
}

方案二：可持久化线段树（主席树） #

思路 #

1. 离散化数组：同方案一，缩小数值范围。
2. 构建可持久化线段树：每个版本的线段树对应原数组前i个元素的统计信息，通过共享节点减少空间开销。
3. 查询中位数：对比版本R和版本L-1的线段树，找到累计元素个数≥k的最小rank，对应数值即为中位数。

该方案查询时间为O(log N)，是效率最优的方案之一，适合对时间要求极高的场景。

方案三：莫队算法 #

思路 #

1. 分块：将数组分成大小为√N的块。
2. 排序查询：按左端点所在块排序，块内按右端点排序（奇偶块交替排序优化）。
3. 维护当前区间：用频率数组或有序集合维护区间元素，动态调整中位数。

时间复杂度为O((N+Q)√N)，实现相对复杂，但对于5e4规模的数据也能通过时间限制。

内容的提问来源于stack exchange，提问作者V K Deewakar