剑指 Offer 40. 最小的k个数

题目

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

题目链接：最小的k个数

示例 1：

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

示例 2：

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

思路

这类问题属于Topk问题，解决这类问题主要有两种思路。

方式1 快排变形

基于快排的变形，基于数组的第k个数字来调整，是小于k的数字都在k的左边（这k个数字不一定是有序的），大于k的数字都在k的右边。空间复杂度o(n),是在修改数组的基础上。

public int[] getLeastNumbers2(int[] arr, int k) {
        int[] ints = new int[k];
        int low=0;
        int high=arr.length-1;
        if(arr==null||k<=0||k>arr.length){
            return ints;
        }
        int index = getIndex(arr, low, high);
        while (index!=k-1){
            if(index>k-1){
                high=index-1;
                index=getIndex(arr,low,high);
            }else{
                low=index+1;
                index=getIndex(arr,low,high);
            }
        }
        for (int i = 0; i <k ; i++) {
            ints[i]=arr[i];
        }
        return ints;
    }
    
    //快排不断调整枢纽的关键代码
    public static int getIndex(int[] arr,int low,int high){
        int temp=arr[low];
        while (low<high){
            while (low<high&&arr[high]>=temp){
                high--;
            }
            arr[low]=arr[high];
            while (low<high&&arr[low]<=temp){
                low++;
            }
            arr[high]=arr[low];
        }
        arr[low]=temp;
        return low;
    }

方式2 使用大顶堆

堆的性质是每次可以找出最大或最小的元素。我们可以使用一个大小为 k 的最大堆（大顶堆）， 将数组中的元素依次入堆，当堆的大小超过 k 时，便将多出的元素从堆顶弹出。由于使用了一个大小为 k 的堆，空间复杂度为 O(k)；入堆和出堆操作的时间复杂度均为 O(logk)，每个元素都需要进行一次入堆操作，故算法的时间复杂度为 O(nlogk)。

public int[] getLeastNumbers(int[] arr, int k){
        if(arr==null||k<=0||k>arr.length){
            return new int[0];
        }
        //PriorityQueue默认是最小堆重写实现最大堆
        Queue<Integer> heap=new PriorityQueue<Integer>(k,(c1,c2)->Integer.compare(c2,c1));

        for (int i : arr) {
            //当前元素小于堆顶元素才入堆
            if(heap.isEmpty()||heap.size()<k||i<heap.peek()){
                heap.offer(i);
            }
            //堆的大小超过 k 时，便将多出的元素从堆顶弹出。
            if(heap.size()>k){
                heap.poll();
            }
        }

        int[] res=new int[heap.size()];
        int j=0;
        for (Integer integer : heap) {
            res[j++]=integer;
        }
        return res;

    }