《剑指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

直接排序法

如果对时间复杂度要求不高情况下，这种方法最直接最快速

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        Arrays.sort(arr);
        int res[] = new int[k];
        for(int i = 0; i < k; i++){
            res[i] = arr[i];
        }
        return res;
    }
}

快速排序法

快排的划分函数每次执行完后都能将数组分成两个部分，小于等于分界值 pivot 的元素的都会被放到数组的左边，大于的都会被放到数组的右边，然后返回分界值的下标。与快速排序不同的是，快速排序会根据分界值的下标递归处理划分的两侧，而这里我们只处理划分的一边。

我们定义函数 randomized_selected(arr, l, r, k) 表示划分数组 arr 的 [l,r] 部分，使前 k 小的数在数组的左侧，在函数里我们调用快排的划分函数，假设划分函数返回的下标是 pos（表示分界值 pivot 最终在数组中的位置），即 pivot 是数组中第 pos - l + 1 小的数，那么一共会有三种情况：

• 如果 pos - l + 1 == k，表示 pivot 就是第 kk 小的数，直接返回即可；
• 如果 pos - l + 1 < k，表示第 kk 小的数在 pivot 的右侧，因此递归调用 randomized_selected(arr, pos + 1, r, k - (pos - l + 1))；
• 如果 pos - l + 1 > k，表示第 kk 小的数在 pivot 的左侧，递归调用 randomized_selected(arr, l, pos - 1, k)。

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

    public void QuickSort(int arr[], int start, int end){
        int pivot = start, i = start, j = end;
        if(start >= end) return;
        while(i < j){
            while (i < j && arr[j] >= arr[pivot]) j--;
            while (i < j && arr[i] <= arr[pivot]) i++;
            int temp = arr[j];
            arr[j] = arr[i];
            arr[i] = temp;
        }

        int temp = arr[pivot];
        arr[pivot] = arr[i];
        arr[i] = temp;

        QuickSort(arr, start, i-1);
        QuickSort(arr, i+1, end);
        return;
    }
}

堆

保持堆的大小为K，然后遍历数组中的数字，遍历的时候做如下判断：

1. 若目前堆的大小小于K，将当前数字放入堆中。
2. 否则判断当前数字与大根堆堆顶元素的大小关系，如果当前数字比大根堆堆顶还大，这个数就直接跳过；

反之如果当前数字比大根堆堆顶小，先poll掉堆顶，再将该数字放入堆中。

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        Queue<Integer> pq = new PriorityQueue<>((v1, v2) -> v2 - v1);
        for (int num: arr) {
            if (pq.size() < k) {
                pq.offer(num);
            } else if (num < pq.peek()) {
                pq.poll();
                pq.offer(num);
            }
        }
        
        // 返回堆中的元素
        int[] res = new int[pq.size()];
        int idx = 0;
        for(int num: pq) {
            res[idx++] = num;
        }
        return res;
    }
}