如何优化寻找逐步扩展子数组第m大元素的Java代码？

问题描述 #

给定一个包含1到n任意排列的n个整数的数组，以及整数m。需依次选取数组的前m个、前m+1个……直至前n个元素作为子数组，分别找出每个子数组中的第m大元素。

示例 #

n = 4
arr = [4,2,1,3]
m = 2

结果：

[2,2,3]

解释：

a) 前m个元素=[4,2]，第2大元素=2
b) 前m+1个元素=[4,2,1]，第2大元素=2
c) 前m+2个元素=[4,2,1,3]，第2大元素=3

原代码的效率问题 #

原代码每次对前i个元素执行逆序排序，再取第m-1位元素。每次排序的时间复杂度为O(i logi)，累加后整体时间复杂度为O(n²logn)，当n较大时，这种方法的效率会急剧下降。

优化方案：使用小顶堆（优先队列） #

核心思路是维护一个大小为m的小顶堆，堆中始终保存当前子数组里最大的m个元素，堆顶元素就是该子数组的第m大元素。这种方法的时间复杂度为O(n logm)，远优于原实现。

具体步骤 #

1. 初始化堆：将数组前m个元素加入小顶堆，此时堆顶就是第一个子数组（前m个元素）的第m大元素，直接加入结果列表。
2. 遍历后续元素：从第m+1个元素开始，逐个处理：
   • 如果当前元素大于堆顶元素，说明它比当前的第m大元素更大，弹出堆顶并将当前元素加入堆，更新堆顶为新的第m大元素。
   • 如果当前元素小于等于堆顶元素，它不会进入前m大的集合，堆顶保持不变。
   • 每次处理后，将堆顶元素加入结果列表。

优化后的Java代码 #

import java.util.ArrayList;
import java.util.List;
import java.util.PriorityQueue;

static List<Integer> solve(List<Integer> arr, int m) {
    List<Integer> result = new ArrayList<>();
    // 初始化大小为m的小顶堆，默认自然顺序，堆顶为最小元素
    PriorityQueue<Integer> minHeap = new PriorityQueue<>(m);

    // 填充前m个元素到堆中
    for (int i = 0; i < m; i++) {
        minHeap.offer(arr.get(i));
    }
    result.add(minHeap.peek());

    // 处理剩余元素
    for (int i = m; i < arr.size(); i++) {
        int currentNum = arr.get(i);
        if (currentNum > minHeap.peek()) {
            minHeap.poll();
            minHeap.offer(currentNum);
        }
        result.add(minHeap.peek());
    }

    return result;
}

代码说明 #

• 小顶堆的大小始终维持在m，确保堆中元素是当前子数组的前m大值，堆顶即为第m大元素。
• 每个元素最多执行一次入堆和出堆操作，每次堆操作的时间复杂度为O(logm)，整体效率大幅提升。

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