题目链接

牛客在线oj题——最小的K个数

题目描述

给定一个长度为 n 的可能有重复值的数组，找出其中不去重的最小的 k 个数。例如数组元素是4,5,1,6,2,7,3,8这8个数字，则最小的4个数字是1,2,3,4(任意顺序皆可)。

数据范围：0≤k,n≤10000，数组中每个数的大0≤val≤1000
要求：空间复杂度 O(n) ，时间复杂度 O(nlogk)

题目示例

示例1

输入：
[4,5,1,6,2,7,3,8],4

返回值：
[1,2,3,4]

说明：
返回最小的4个数即可，返回[1,3,2,4]也可以

示例2

输入：
[1],0

返回值：
[]

示例3

输入：
[0,1,2,1,2],3

返回值：
[0,1,1]

解题思路

这道题是一个典型的topK问题，可以直接使用任意一种排序，将数组中的元素排成从小到大的顺序，然后直接返回前k个数字

而更好的解法是，构造一个含k个元素的大根堆，先讲数组的前k个元素插入到堆中，然后继续从左到右遍历数组中的元素

大根堆的特性是：堆顶元素比下面所有元素都要大

当遍历到的元素i比大根堆堆顶的元素peek大时，说明i肯定不是最小的k个数之中的元素，继续遍历下一个位置

当遍历到的元素i比大根堆堆顶的元素peek小时，说明peek肯定不是最小的k个数中的元素，把大根堆堆顶元素弹出，将i插入到大根堆中

遍历完整个数组序列，最终大根堆中所有元素就是最小的k个数

例如：

首先将前四个元素插入大根堆

当前遍历到的元素是2，比大根堆堆顶元素6小，将6弹出，将2入大根堆，i++

当前遍历到的元素是7，比大根堆堆顶元素5大，肯定不在最小的k个数中，继续遍历

当前遍历到的元素为3，比大根堆堆顶元素5小，将5弹出大根堆，将3插入到大根堆中

当前遍历到的元素是8，比大根堆堆顶元素4大，肯定不在最小的k个数中，遍历结束，当前大根堆中存放的所有元素就是最小的k个数

完整代码

import java.util.*;

public class Solution {
    public ArrayList<Integer> GetLeastNumbers_Solution(int [] input, int k) {
        ArrayList<Integer> result = new ArrayList<>();
        if(input == null || k <= 0 || k > input.length){
            return result;
        }
        PriorityQueue<Integer> priorityQueue = new PriorityQueue<>(k, Collections.reverseOrder());
        for (int i = 0; i < input.length; i++){
            if(i < k){
                priorityQueue.offer(input[i]);
            } else {
                if(input[i] < priorityQueue.peek()){
                    priorityQueue.poll();
                    priorityQueue.offer(input[i]);
                }
            }
        }
        for (int i = 0; i < k; i++){
            result.add(priorityQueue.poll());
        }
        return result;
    }
}