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

解题思路

第一种解法

1、使用Partition函数来解决这个问题。
如果基于数组的第k个数字来调整，则使得比第k个数字小的所有数字都位于数组的左边，比第k个数字大的所有数字都位于数组的右边。这样调整之后，位于数组总左边的k个数字就是最小的k个数字（这k个数字不一定是排序的）
2、需要注意的一个细节，边界问题，没有考虑到的话会报错

if(k == 0){
	return new int[0];
}else if(arr.length <= k){
	return arr;
}

3、时间复杂度为O(n)，但是需要修改输入的数组，使用Partition会调整数组中数字的顺序。

代码

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if(k == 0){
            return new int[0];
        }else if(arr.length <= k){
            return arr;
        }

        quickselect(arr, 0, arr.length - 1, k);
        return Arrays.copyOfRange(arr, 0, k);
    }

    public void quickselect(int[] nums, int start, int end, int k){
        int pivot = nums[end];
        int left = start;

        for(int i = start; i < end; i++){
            if(nums[i] <= pivot)
                swap(nums, left++, i);
        }
        swap(nums, left, end);

        if(left == k)
            return;
        else if(left < k)
            quickselect(nums, left + 1, end, k);
        else
            quickselect(nums, start,  left - 1, k);
    }
    void swap(int[] nums, int i, int j){
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}

第二种解法

大根堆求前k小Java中提供的现成的PriorityQueue默认是小根堆
先创建一个容器来存储k个数字。找出已有k个数中的最大值，然后拿这次待插入的整数和最大值进行比较。如果待插入的值比当前已有的最大值小，则用这个数替换当前已有的最大值，如果待插入的值比当前已有的最大值还要大，则抛弃这个数。

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;

    }
}