剑指Offer 40. 最小的k个数

摘要：
经典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$

---

快速查找方法(分治法，快速排序思想)

• Java

class TopK {
    public static int[] getTopK(int[] arr, int k) {
        if (k == 0)
            return new int[0];
        else if (k >= arr.length)
            return arr;
        int low = 0;
        int high = arr.length -1;
        partitionArray(arr, low, high, k);
        //copy array
        int[] result = new int[k];
        for (int copy = 0; copy < k; copy++){
            result[copy] = arr[copy];
        }
        return result;
    }

    private static void partitionArray(int[] arr, int low, int high, int k) {
        int m = partition(arr, low, high);
        if (k == m) //as we expect，反复递归，k会在m左右两侧摇摆，直到恰好等于m
            return;
        else if (k < m)
            partitionArray(arr, low, m -1, k);
        else
            partitionArray(arr, m + 1, high, k);
    }

    private static int partition(int[] arr, int low, int high){
        int i = low;
        int j = high + 1;
        int v = arr[low];

        while (true) {
            while (arr[++i] < v) {
                if (i == high) {
                    break;
                }
            }
            while (arr[--j] > v) {
                if (j == low) {
                    break;
                }
            }

            if (i >= j) { //碰头了
                break;
            }
            swap(arr, i, j);
        }
        swap(arr, low, j);
        return j;
    }

    private static void swap(int[]arr, int i, int j){
        int t = arr[i];
        arr[i] = arr[j];
        arr[j] = t;
    }

堆方法

维护一个大顶堆，堆中始终保持当前状态最小k个数。

• C++
• Java

class Solution {
public:
    priority_queue<int> mxheap;
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        for (int i = 0; i < arr.size(); i++) {
            mxheap.push(arr[i]);
            for (; mxheap.size() > k; mxheap.pop());
        }
        vector<int> res;
        for (; !mxheap.empty(); mxheap.pop())
            res.push_back(mxheap.top());
        return res;
    }
};

class Solution {
    static public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0) {
            return new int[0];
        }
        // 使用一个最大堆（大顶堆）
        // Java 的 PriorityQueue 默认是小顶堆，添加 comparator 参数使其变成最大堆,or we can do it like this:
        // Collections.reverseOrder()
        //PriorityQueue<Integer> heap = new PriorityQueue<>(Collections.reverseOrder());
        PriorityQueue<Integer> heap = new PriorityQueue<>(k, (i1, i2) -> Integer.compare(i2, i1));

        for (int e : arr) {
            // 当前数字小于堆顶元素才会入堆
            if (heap.isEmpty() || heap.size() < k || e < heap.peek()) { //Queue.peek()返回堆顶最大元素（不删除），空堆返回null
                heap.offer(e); //Queue.offer()插入元素e，返回true
            }
            if (heap.size() > k) {
                heap.poll(); // Queue.poll()推出堆顶最大元素作为返回值，空堆返回null
            }
        }
        // 将堆中的元素存入数组
        int[] res = new int[heap.size()];
        int j = 0;
        for (int e : heap) { //遍历堆
            res[j++] = e;
        }
        return res;
    }
}

二叉搜索树

待更新。

原题链接： 剑指Offer 40. 最小的k个数