剑指 Offer 40. 最小的k个数

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

剑指 Offer 40. 最小的k个数 - 力扣（LeetCode） (leetcode-cn.com)

//利用快速排序的思想，随机选取一个数字作为基准，使得比它小的数字位于它的左边，比它大的数字位于它的右边
//如果基准数字正好位于第k个位置，即下标为k-1，则左边的k个数字就是最小的k个数
//如果基准数字的下标大于k-1，就对左边部分递归，直到找到位于第k个位置的数字
//如果基准数字的下标小于k-1，就对右边部分递归
class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if(arr == null || arr.length == 0 || k == 0){
            return new int[0];
        }
        partion(arr, 0, arr.length - 1, k);
        int[] res = new int[k];
        for(int i = 0; i < k; i++){
            res[i] = arr[i];
        }
        return res;
    }

    public void partion(int[] arr, int start, int end, int k){
        int rand = (int)(Math.random() * (end - start + 1)) + start;
        int p = start;
        swap(arr, rand, end);
        for(int i = start; i < end; i++){
            if(arr[i] <= arr[end]){
                if(p != i){
                    swap(arr, p, i);
                }
                p++;
            }
        }
        swap(arr, p, end);
        if(p == k - 1){
            return;
        }else if(p < k){
            partion(arr, p + 1, end, k);
        }else{
            partion(arr, start, p - 1, k);
        }
    }

    public void swap(int[] arr, int i, int j){
        int temp = arr[i];
        arr[i] = arr[j];
        arr[j] = temp;
    }
}

//利用大小为k的大顶堆存储数据，首先把k个数字放入大顶堆，最大的数字就位于堆顶，
//接下来遍历剩余的数字，如果当前数字小于堆顶数字，就用把堆顶数字替换为当前数字，
//遍历完成后，大顶堆里存储的就是最小的k个数
class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if(arr == null || arr.length == 0 || k == 0){
            return new int[0];
        }
        PriorityQueue<Integer> queue = new PriorityQueue(new Comparator<Integer>(){
            public int compare(Integer num1, Integer num2){
                return -(num1 - num2);
            }
        });
        for(int i = 0; i < k; i++){
            queue.offer(arr[i]);
        }

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

        int[] res = new int[k];
        for(int i = 0; i < k; i++){
            res[i] = queue.poll();
        }

        return res;
    }

}