剑指 Offer 40. 最小的k个数

题目

输入整数数组 arr ，找出其中最小的 k 个数。例如，输入4、5、1、6、2、7、3、8这8个数字，则最小的4个数字是1、2、3、4。

示例 1：

输入：arr = [3,2,1], k = 2
输出：[1,2] 或者 [2,1]

示例 2：

输入：arr = [0,1,2,1], k = 1
输出：[0]

限制：

• 0 <= k <= arr.length <= 10000
• 0 <= arr[i] <= 10000

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/zui-xiao-de-kge-shu-lcof

解题思路

Top K类问题,两种思路:快速排序和堆排序

这里解释一下堆排序,求最小的K个数

• 维护一个大小为k的大顶堆
• 当堆的大小没有达到k时,元素直接加入堆中
• 当堆的大小等于k时,最大的元素在堆顶,当前元素小于堆顶元素的话,堆顶元素出堆,当前元素入堆

java中的PriorityQueue是一个优先队列,默认就是小顶堆,可以覆写比较器成大顶堆

代码

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        // 特殊情况判断
        if (arr == null || arr.length == 0 || k == 0) {
            return new int[0];
        }
        
        // 使用java的优先队列,默认是小顶堆,复写构建大顶堆
        PriorityQueue<Integer> queue = new PriorityQueue<>((v1, v2) -> v2 - v1);
        for (int item : arr) {
            if (queue.size() < k) {
                // 先把堆填充满
                queue.offer(item);
            } else if (item < queue.peek()) {
                // 比较当前元素和堆顶的大小,小于堆顶元素则堆顶元素出队,当前元素入队
                queue.poll();
                queue.offer(item);
            }
        }

        // 把队列中的元素构建成数组返回
        int[] result = new int[queue.size()];
        for (int i = 0; i < result.length; i++) {
            result[i] = queue.poll();
        }
        return result;
    }
}

复杂度

• 时间复杂度:O(nlogk),入堆和出堆的时间复杂度为O(logk),极端情况下每个元素执行一次入堆操作,即最差情况下,时间复杂度为O(nlogk)
• 空间复杂度:O(k),维护了一个大小为k的堆