算法-TopK问题

描述

输入整数数组 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

输入：{1,2,3}

返回值：{3,2,1}

解决方案

方案一：利用大根堆

java默认的是小根堆，我们用一个大根堆实时维护数组的前 k 小值。首先将前 k 个数插入大根堆中，随后从第 k+1
个数开始遍历，如果当前遍历到的数比大根堆的堆顶的数要小，就把堆顶的数弹出，再插入当前遍历到的数。最后将大根堆里的数存入数组返回即可。

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (arr == null || arr.length == 0 || k < 1) {
            return new int[0];
        }
        if (k >= arr.length) {
            return arr.clone();
        }
        PriorityQueue<Integer> bigEndQueue = new PriorityQueue<>(k, Comparator.reverseOrder());
        for (int i = 0; i < k; i++) {
            bigEndQueue.offer(arr[i]);
        }
        for (int i = k; i < arr.length; i++) {
            if (arr[i] < bigEndQueue.peek()) {
                bigEndQueue.poll();
                bigEndQueue.offer(arr[i]);
            }
        }
        int[] res = new int[k];
        for (int i = 0; i < k; i++) {
            res[i] = bigEndQueue.poll();
        }
        return res;
    }
}

方案二：利用partition的思想

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (arr == null || arr.length == 0 || k <= 0) {
            return new int[0];
        }
        if (k >= arr.length) {
            return arr.clone();
        }
        getLeastNumbers0(arr, 0, arr.length, k);
        int[] res = new int[k];
        for (int i = 0; i < k; i++) {
            res[i] = arr[i];
        }
        return res;
    }

    public void getLeastNumbers0(int[] arr, int l, int r, int k) {
        if (l >= r) {
            return;
        }

        int pivot = arr[r - 1];
        int i = l;
        int pos = l - 1;
        while (i < r) {
            if (arr[i] <= pivot) {
                swap(arr, ++pos, i);
            }
            i++;
        }
        if (pos == k) {
            return;
        }
        if (pos < k) {
            getLeastNumbers0(arr, pos + 1, r, k);
        } else {
            getLeastNumbers0(arr, l, pos, k);
        }
    }

    private void swap(int[] nums, int i, int j) {
        if (i == j) {
            return;
        }
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }
}