当前位置：首页 > news >正文

Java--利用(堆)获取前k个最小元素

news 2025/5/18 8:09:42

一、问题背景

二、核心方法解析

关键点说明：

三、优化思路（进阶）

四、完整代码实现

五、方法对比

六、面试题

七、总结

正文开始

一、问题背景

在处理数据时，我们经常需要快速获取一个数组中前k个最小的元素。例如：

在推荐系统中筛选评分最低的商品
在数据分析中提取温度最低的k个城市

本文将介绍如何通过Java的PriorityQueue高效实现这一功能。

二、核心方法解析

// 方法签名
public int[] smallestK(int[] arr, int k) {PriorityQueue<Integer> priorityQueue = new PriorityQueue<>();for (int num : arr) {priorityQueue.offer(num); // 所有元素入队（默认小根堆）}int[] ret = new int[k];for (int i = 0; i < k; i++) {ret[i] = priorityQueue.poll(); // 依次取出堆顶元素（即最小值）}return ret;
}

关键点说明：

优先级队列特性：Java的PriorityQueue默认是小根堆，堆顶始终为最小元素
时间复杂度：O(n log n) + O(k log n) → 总体为O(n log n)
空间复杂度：O(n)（需存储所有元素）

三、优化思路（进阶）

当数据量极大时，可采用大根堆优化，将时间复杂度降低至O(n log k)：

// 使用大根堆（仅保留k个最小元素）
PriorityQueue<Integer> maxHeap = new PriorityQueue<>((a, b) -> b - a);
for (int num : arr) {if (maxHeap.size() < k) {maxHeap.offer(num);} else if (num < maxHeap.peek()) {maxHeap.poll();maxHeap.offer(num);}
}

四、完整代码实现

import java.util.PriorityQueue;public class TopKSmallest {// 方法1：直接使用小根堆public int[] smallestK(int[] arr, int k) {if (k <= 0) return new int[0];PriorityQueue<Integer> pq = new PriorityQueue<>();for (int num : arr) {pq.offer(num);}int[] result = new int[k];for (int i = 0; i < k; i++) {result[i] = pq.poll();}return result;}// 方法2：大根堆优化版（推荐）public int[] smallestKOptimized(int[] arr, int k) {if (k <= 0) return new int[0];PriorityQueue<Integer> maxHeap = new PriorityQueue<>((a, b) -> b - a);for (int num : arr) {if (maxHeap.size() < k) {maxHeap.offer(num);} else if (num < maxHeap.peek()) {maxHeap.poll();maxHeap.offer(num);}}return maxHeap.stream().mapToInt(i -> i).toArray();}public static void main(String[] args) {int[] data = {12, 3, 5, 7, 19, 2, 16};TopKSmallest solver = new TopKSmallest();System.out.println(Arrays.toString(solver.smallestK(data, 3)));    // [2, 3, 5]System.out.println(Arrays.toString(solver.smallestKOptimized(data, 3))); // [5, 3, 2]}
}