【大厂面试真题-字节】数组中的前k个数

题目描述

设计一个算法，找出数组中最小的k个数。以任意顺序返回这k个数均可。

示例

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

思路

首先分析题意，从数组中找到最小的k的数，最先想到的就是将数组进行由大到小排序，输出前k个数值，这是最容易想到的办法。但是对于这种topk问题，更加常用的解法是使用优先队列。这里我们可以采用大顶堆的方式，也就是最顶部永远是堆中最大的数。我们初始化一个k大小的大顶堆，将数组的前k个数入堆，然后向后遍历数组，将每一个数值与堆顶元素进行比较，如果当前元素小于堆顶元素，则证明堆中元素不是前k小的数，将堆顶元素出堆，当前元素入堆，依次遍历完数组之后，堆中存储的元素即为数组中前k个元素。

实现

public int[] smallestK(int[] arr, int k) {
        //1。初始化一个大顶堆
        PriorityQueue<Integer> queue = new PriorityQueue<>((o1, o2) -> o2 - o1);
        //2。将arr中的前k个元素入堆
        for (int i = 0 ; i < k ; i ++) {
            queue.add(arr[i]);
        }
        //3.开始数组后续元素的比较，如果小于堆顶元素，堆定元素出来，当前元素进去
        for (int i = k ; i < arr.length ; i++) {
            if (arr[i] < queue.peek()) {
                queue.poll();
                queue.add(arr[i]);
            }
        }
        //4.堆中存在的元素应该就是arr的前k小的数值
        int[] result = new int[k];
        int index = 0;
        while (!queue.isEmpty()) {
            result[index] = queue.poll();
            ++ index;
        }
        return result;
    }