剑指 Offer 40. 最小的k个数

题目描述

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

输入输出

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

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

基本思路

用排序算法做

快排：时间复杂度$O(n)$ 空间复杂度$O(logn)$

堆排序：时间复杂度$O(nlogk)$ 空间复杂度$O(k)$ Java中堆排序可以使用优先队列来实现

Reference：PriorityQueue

java实现

// 直接调库 面试直接寄
class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        int[] vec = new int[k];
        Arrays.sort(arr);
        for (int i = 0; i < k; ++i) {
            vec[i] = arr[i];
        }
        return vec;
    }
}

//快排
class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        int[] res = new int[k];
        QuickSort(arr, 0, arr.length - 1);
        for(int i = 0; i < k; i++){
            res[i] = arr[i];
        }
        return res;
    }

    void QuickSort(int[] arr, int low, int high){
        if(low < high){
            int pivot = Partition(arr, low, high);
            QuickSort(arr, low, pivot-1);
            QuickSort(arr, pivot+1, high);
        }
        
    }
    
    int Partition(int[] arr, int low, int high){
        int pivot = arr[low];
        while(low < high){
            while(low < high && arr[high] >= pivot) high--;
            arr[low] = arr[high];
            while(low < high && arr[low] <= pivot) low++;
            arr[high] = arr[low];
        }
        arr[low] = pivot;
        return low;
    }
}

//堆排序
class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        int[] res = new int[k];
        if(k == 0 || arr.length == 0){
            return res;
        }
        
        PriorityQueue<Integer> pq = new PriorityQueue<Integer>(new Comparator<Integer>(){
            public int compare(Integer num1, Integer num2){
                return num2.compareTo(num1);//大顶堆 队首元素最大
            }
        });

        for(int i = 0; i < k; i++){
            pq.offer(arr[i]);
        }
        for(int i = k; i < arr.length; i++){
            if(pq.peek() > arr[i]){
                pq.poll();
                pq.offer(arr[i]);
            }
        }
        for(int i = 0; i < k; i++){
            res[i] = pq.poll();
        }
        return res;
    }
}