【剑指offer】53. 最小的k个数（java选手）（优先队列+快排+快速选择）

题目链接

acwing题目链接
力扣题目链接

题目描述

输入 n个整数，找出其中最小的 k 个数。

注意：
输出数组内元素请按从小到大顺序排序;
数据范围
1≤k≤n≤1000

样例
输入：[1,2,3,4,5,6,7,8] , k=4
输出：[1,2,3,4]

题目分析

• 排序算法

题解

优先队列-小根堆

• 用小根堆

将所有元素都放入小根堆中，就相当于堆元素进行了排序。
然后依次从优先队列的首部（最小的位置）开始取元素，取k个为止

注意啦！
这是优先队列（升序排列）实现的小根堆
也就是每次peek取出的元素都是当前优先队列中最小的！

class Solution {
    public List<Integer> getLeastNumbers_Solution(int[] input, int k) {
        // 优先队列（升序），小根堆，根是最小的【peek取出的元素是最小的】
        PriorityQueue<Integer> pq = new PriorityQueue<>((a, b)->a-b);
        for(int i = 0; i < input.length; i ++){
            pq.add(input[i]);
        }
        
        List<Integer> list = new ArrayList<>();
        for(int i = 0; i < k; i ++){
            list.add(pq.poll());
        }
  
        return list;
    }
}

优先队列-大根堆

• 从大到小排序，首部元素是最大的
• 只存放k个整数，如果超过这个值之后，遍历的元素与首部元素小才插入（因为要找前k个最小的元素）
• 遍历优先队列，不停的在list首部插入值（因为是大根堆，所以队列首部的元素是最大的）

class Solution {
    public List<Integer> getLeastNumbers_Solution(int[] input, int k) {
        // 大根堆，堆中只存放k个整数，超过这个k值后，如果元素比peek得到的值小，插入
        // 优先队列从大到小排序，逆序
        PriorityQueue<Integer> pq = new PriorityQueue<>((a, b)-> b - a);
        
        for(int i = 0; i < input.length; i ++){
            if(pq.size() >= k){
                if(pq.peek() > input[i]){
                    pq.poll();
                    pq.add(input[i]);
                }
            }else{
                pq.add(input[i]);
            }
        }
        
        List<Integer> list = new ArrayList<>();
        for(int i = 0; i < k; i ++){
            list.add(0, pq.poll());
        }
        
        return list;
    }
}

！0323更新！上述代码在力扣上无法通过，没有考虑到k为0的情况；下面是更新版本

class Solution {
    public int[] smallestK(int[] arr, int k) {
        if(k==0) return new int[0];
        // 优先队列，降序
        PriorityQueue<Integer> pq = new PriorityQueue<>((a,b)-> b-a);
        for(int i = 0; i < arr.length; i ++){
            if(pq.size() >= k){
                if(arr[i] < pq.peek()){
                    pq.poll();
                    pq.add(arr[i]);
                }  
            }else{
                pq.add(arr[i]);
            }
        }
        int[] res = new int[k];
        for(int i = 0; i < k; i ++){
            res[i] = pq.poll();
        }
        return res;
    }
}

快排

写法1：

class Solution {
    public int[] smallestK(int[] arr, int k) {
        sort(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 sort(int[] arr, int left, int right){
        // 只有最后一个元素
        if(left >= right) return;

        // 基准
        int base = arr[left];
        int i = left, j = right;
        while(i < j){
            while(i < j && arr[j] > base){
                j--;
            }
            if(i < j){
                arr[i] = arr[j];
                i ++;
            }


            while(i < j && arr[i] < base){
                i ++;
            }
            if(i < j){
                arr[j] = arr[i];
                j --;
            }
        }

        arr[i] = base;
        sort(arr, left, i - 1);
        sort(arr, i + 1, right);
    }
}

写法2：

class Solution {
    public int[] smallestK(int[] arr, int k) {
        sort(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 sort(int[] arr, int left, int right){
        // 只有最后一个元素
        if(left >= right) return;

        // 基准
        int base = arr[(left+right)/2];
        int i = left-1, j = right+1;
        while(i < j){
           do i++; while(arr[i] < base);
           do j--; while(arr[j] > base);
           if(i < j){
            int temp = arr[i];
            arr[i] = arr[j];
            arr[j] = temp;
           }
        }
        sort(arr, left, j);
        sort(arr, j + 1, right);
    }
}

快速选择

题目链接；找第k个元素

class Solution {
    public int findKthLargest(int[] nums, int k) {
        // 需要数组中下标为need的元素
        int need = nums.length - k;
        return sort(nums, 0, nums.length - 1, need);
    }

    public int sort(int[] nums, int left, int right, int k) {
        if(left == right) return nums[left];

        int base = nums[(left + right) / 2];
        int i = left - 1, j = right + 1;
        while (i < j) {
            do i++; while (nums[i] < base);
            do j--; while (nums[j] > base);
            if (i < j) {
                int temp = nums[i];
                nums[i] = nums[j];
                nums[j] = temp;
            }
        }

        if(j >= k) return sort(nums, left, j, k);
        else return sort(nums, j + 1, right, k);
    }
}

注意！边界

if(j >= k) return sort(nums, left, j, k);
else return sort(nums, j + 1, right, k);