求数组中每个长度为m的连续子数组的第k小元素

给定一个长度为n的数值数组，选取长度为m的连续子数组，找出每个子数组中的第k小元素。

示例：

arr = [1, 3, 2, 1], m = 3, k = 2
m表示子数组长度
k表示子数组中的第k小元素
即找出长度为3的子数组中的第2小元素。

共有2个连续子数组：
[1, 3, 2] 和 [3, 2, 1]

k = 2，上述子数组的第2小元素为 [2,2]

约束条件：

k, m , n 的取值范围为1到10^5
数组元素取值范围为1到10^9

我尝试的两种解法： #

方法1：暴力法 #

这是朴素解法：

public static int[] solve(int[] arr, int m, int k) {
    List<Integer> result = new ArrayList<>();
    int n = arr.length;
    for(int i=0; i<=n-m; i++) {
       int[] sub = Arrays.copyOfRange(arr, i, i+m);
       Arrays.sort(sub);
       result.add(sub[k-1]);
    }
    return result;
}

该方法时间效率极低，每次排序子数组需O(m log m)，总复杂度O(n m log m)，完全无法处理1e5规模的输入。

方法2：优先队列法 #

public static int[] solve(int[] arr, int m, int k) {
    int n = arr.length;
    int[] result = new int[n - m + 1];
    Queue<Integer> q = new PriorityQueue<>((a, b) -> Integer.compare(b, a));

    for (int i = 0; i < m; i++) {
        q.add(arr[i]);
        if (q.size() > k) {
            q.poll();
        }
    }
    result[0] = q.peek();

    for (int i = 1; i <= n - m; i++) {
        int left = arr[i - 1];
        if (q.contains(left)) {
            q.remove(left);
        } else {
            q.poll();
        }
        q.add(arr[i+m-1]);
        result[i] = q.peek();
    }

    return result;
}

该方法存在两个致命问题：

1. PriorityQueue.contains()和remove()是O(k)复杂度，当k接近1e5时，单次操作耗时极高，直接超时。
2. 逻辑错误：当窗口移出的元素不在堆中时，直接poll()会错误减少堆的大小。堆中维护的是窗口内前k小的元素，移出的元素可能是比堆顶大的元素（不在堆里），此时堆的大小仍为k，不需要任何删除操作，否则后续添加新元素后堆顶不再是第k小元素。

正确解法 #

解法1：有序计数集合维护滑动窗口 #

利用TreeMap（基于红黑树）统计窗口内元素的出现次数，通过遍历有序的键快速找到第k小元素，所有操作的时间复杂度为O(log m)。

import java.util.TreeMap;

public static int[] solve(int[] arr, int m, int k) {
    int n = arr.length;
    int[] result = new int[n - m + 1];
    TreeMap<Integer, Integer> countMap = new TreeMap<>();

    // 初始化第一个窗口的元素计数
    for (int i = 0; i < m; i++) {
        countMap.put(arr[i], countMap.getOrDefault(arr[i], 0) + 1);
    }
    result[0] = findKthSmallest(countMap, k);

    // 滑动窗口处理后续子数组
    for (int i = 1; i <= n - m; i++) {
        // 移除窗口左端元素
        int leftVal = arr[i - 1];
        countMap.put(leftVal, countMap.get(leftVal) - 1);
        if (countMap.get(leftVal) == 0) {
            countMap.remove(leftVal);
        }
        // 添加窗口右端新元素
        int rightVal = arr[i + m - 1];
        countMap.put(rightVal, countMap.getOrDefault(rightVal, 0) + 1);
        // 找到当前窗口的第k小元素
        result[i] = findKthSmallest(countMap, k);
    }

    return result;
}

// 遍历有序集合，累加计数找到第k小元素
private static int findKthSmallest(TreeMap<Integer, Integer> countMap, int k) {
    int sum = 0;
    for (int num : countMap.keySet()) {
        sum += countMap.get(num);
        if (sum >= k) {
            return num;
        }
    }
    return -1; // 因k<=m<=n，不会执行到此处
}

复杂度分析：每个元素的添加、删除操作是O(log m)，每次查找第k小元素最多遍历到第k个不同元素（实际因TreeMap有序，平均复杂度更低），总时间复杂度O(n log m)，可以轻松处理1e5规模的输入。

解法2：二分查找+前缀和离线处理 #

适合元素范围极大但数组规模大的场景，时间复杂度O(n log(max_val))：

1. 对元素值进行二分，判断某个值x是否满足「窗口内<=x的元素数量>=k」。
2. 用前缀和数组快速计算任意窗口内<=x的元素数量。
3. 对每个窗口，找到满足条件的最小x，即为该窗口的第k小元素。

实现略复杂，但效率更高，适合极端场景。

内容的提问来源于stack exchange，提问作者Learner