剑指 Offer 40. 最小的k个数

题目

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

思路

TopK问题两种经典解法（面试常问）

1. 快速选择（快速排序变体）
2. 堆（优先队列）

类似LeetCode HOT 100 —— 215.数组中的第K个最大元素

方法一：快速选择

快速排序中有一步很重要的操作是 partition（划分），从数组中随机选取一个枢纽元素 pivot，然后原地移动数组中的元素，使得比 pivot 小的元素在 pivot 的左边，比 pivot 大的元素在 pivot 的右边

这个 partition 操作是原地进行的，需要 O(n) 的时间，接下来，快速排序会递归地排序左右两侧的数组。而快速选择（quick select）算法的不同之处在于，接下来只需要递归地选择一侧的数组。快速选择算法想当于一个“不完全”的快速排序，因为我们只需要知道最小的 k 个数是哪些，并不需要知道它们的顺序。

目的是寻找最小的 k 个数。假设经过一次 partition 操作，枢纽元素位于下标 m，也就是说，左侧的数组有 m 个元素，是原数组中最小的 m 个数

• 若 k=m，就找到了最小的 k 个数，就是左侧的数组；
• 若 k<m ，则最小的 k 个数一定都在左侧数组中，只需要对左侧数组递归地 parition 即可；
• 若 k>m，则左侧数组中的 m 个数都属于最小的 k 个数，还需要在右侧数组中寻找最小的 k−m个数，对右侧数组递归地 partition 即可

java代码如下：

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        // 最后一个参数表示要找的是下标为k-1的数
        return quickSearch(arr, 0, arr.length - 1, k - 1);
    }

    private int[] quickSearch(int[] nums, int l, int r, int k) {
        // 每快排切分1次，找到排序后下标为j的元素，如果j恰好等于k就返回j以及j左边所有的数；
        int j = partition(nums, l, r);
        if (j == k) {
            return Arrays.copyOf(nums, j + 1);
        }
        // 否则根据下标j与k的大小关系来决定继续切分左段还是右段。
        return j > k ? quickSearch(nums, l, j - 1, k): quickSearch(nums, j + 1, r, k);
    }

    // 快排切分，返回下标j，使得比nums[j]小的数都在j的左边，比nums[j]大的数都在j的右边。
    private int partition(int[] nums, int l, int r) {
        int pivot = nums[l];
        int i = l, j = r + 1;
        while (true) {
            while (++i <= r && nums[i] < pivot);
            while (--j >= l && nums[j] > pivot);
            if (i >= j) {
                break;
            }
            int t = nums[j];
            nums[j] = nums[i];
            nums[i] = t;
        }
        nums[l] = nums[j];
        nums[j] = pivot;
        return j;
    }
}

方法二：堆（优先队列）

堆的性质是每次可以找出最大或最小的元素。

可以使用一个大小为 k 的大顶堆（第/前k小用大顶堆），将数组中的元素依次入堆，当堆的大小超过 k 时，便将多出的元素从堆顶弹出。

java代码如下：

// 保持堆的大小为K，然后遍历数组中的数字，遍历的时候做如下判断：
// 1. 若目前堆的大小小于K，将当前数字放入堆中。
// 2. 否则判断当前数字与大根堆堆顶元素的大小关系，如果当前数字比大根堆堆顶还大，这个数就直接跳过；
//    反之如果当前数字比大根堆堆顶小，先poll掉堆顶，再将该数字放入堆中。
class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        // java中优先队列默认是小根堆，实现大根堆需要重写一下比较器。
        Queue<Integer> pq = new PriorityQueue<>((v1, v2) -> v2 - v1);//注意，定义大根堆要降序定义（因为大根堆的根节点最大，要保证降序），要区别出 大根堆用于升序排列（因为大根堆每次出的根节点是最大的，所以用于升序排列）
        for (int num: arr) {
            if (pq.size() < k) {//如果堆的容量小于k
                pq.offer(num);//加入堆
            } else if (num < pq.peek()) {//如果堆的容量大于等于k，且要加入的元素num比堆顶元素小，则出堆顶元素，加入num
                pq.poll();
                pq.offer(num);
            }
        }
        
        // 返回堆中的元素
        int[] res = new int[pq.size()];
        int idx = 0;
        for(int num: pq) {
            res[idx++] = num;
        }
        return res;
    }
}

两种方法的优劣（面试常问）

快速选择 相比 堆，快速选择 的时间复杂度和空间复杂度 都优于使用堆的方法

但是使用快速选择有一些局限性：

1. 算法需要修改原数组，如果原数组不能修改的话，还需要拷贝一份数组，空间复杂度就上去了
2. 算法需要保存所有的数据。如果把数据看成输入流的话，使用堆的方法是来一个处理一个，不需要保存数据，只需要保存 k 个元素的最大堆。而快速选择的方法需要先保存下来所有的数据，再运行算法。当数据量非常大的时候，甚至内存都放不下的时候，就麻烦了。所以当数据量大的时候还是用基于堆的方法比较好