剑指 Offer 40. 最小的k个数

题目

思路

• 先用快速选择找到第K小的数
• 然后再扫原数组添加即可。

快速选择

注意在添加res时的几种特殊情况

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0) return new int[0];
        int length = arr.length;
        int[] res = new int[k];
        findKthSmallNumber(arr, 0, length - 1, k);
        int kthNum = arr[k - 1];//得到第K小的数
        for (int i = 0, j = 0; i < length && j < k; i++) {// j<k
            if (arr[i] <= kthNum) {//一定是<=，处理1,2,2,2并且kthnum=2的情况
                res[j] = arr[i];
                j++;
            }
        }
        //结果集的最后一个数，一定是kthNum,因为数组可能有重复
        // 如：1,1,1,1,1,2,2,2,2,3,3, k = 2
        res[k - 1] = kthNum;
        return res;
    }

    //快速查找把第K小的数换到arr中索引为k-1的地方
    public void findKthSmallNumber(int[] arr, int start, int end, int k) {
        int i = start, j = end;
        int pivot = start;
        while (i != j) {
            while (arr[i] < arr[pivot]) {
                i++;
            }
            while (i < j && arr[j] >= arr[pivot]) {
                j--;
            }
            swap(arr, i, j);
        }
        swap(arr, pivot, i);
        if (i == k - 1) {
            return;
        } else if (i > k - 1) {
            findKthSmallNumber(arr, start, i - 1, k);
        } else {
            findKthSmallNumber(arr, i + 1, end, k);
        }
    }

    public void swap(int[] arr, int m, int n) {
        int tmp = arr[m];
        arr[m] = arr[n];
        arr[n] = tmp;
    }
    }

大顶堆

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0) {
            return new int[0];
        }
        Queue<Integer> pq = new PriorityQueue<>((v1, v2) -> v2 - v1);
        for (int x : arr) {
            if (pq.size() < k) {
                pq.offer(x);
            } else {
                if (x < pq.peek()) {
                    pq.poll();
                    pq.offer(x);
                }
            }
        }
        int[] res = new int[k];
        for (int i = 0; i < k; i++) {
            res[i] = pq.poll();
        }
        return res; 
    }
}