剑指 Offer 40. 最小的 k 个数

参考题解：4种解法秒杀 TopK (快排、堆、二叉搜索树 BST、计数排序)

这道题基于排序做十分简单。但是，由于只需要求最小的 k 个数，因此其实不用全部排序，只需要排序最小的 k 个即可。

因此，考虑做法：

1. 快速排序

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        return quickSort(arr, 0, arr.length - 1, k - 1);
    }

    public int[] quickSort(int[] nums, int l, int h, int k) {
        int j = partition(nums, l, h);
        if (j == k) {
            return Arrays.copyOf(nums, j + 1);
        }
        return j > k ? quickSort(nums, l, j - 1, k) : quickSort(nums, j + 1, h, k);
    }

    public int partition(int[] nums, int l, int h) {
        int v = nums[l];
        int i = l, j = h + 1;
        while (true) {
            while (++i <= h && nums[i] < v);
            while (--j >= l && nums[j] > v);
            if (i >= j) {
                break;
            }
            int t = nums[j];
            nums[j] = nums[i];
            nums[i] = t;
        }
        nums[l] = nums[j];
        nums[j] = v;
        return j;
    }
}

2. 堆

由于要求前 k 小，因此选择大根堆，每次把最大的 poll 出，从而让堆中元素最小。

这里是很有趣的逆向思维。一般我们要堆排序求最小，都是用小根堆。但那样的话，复杂度就是 O(nlogn)，而不是 O(nlogk) 了。这里的逆向思维很妙。

基于 Java 提供的 PriorityQueue 可以快速实现。

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

3. 二叉搜索树 BST

二叉搜索树求得的前 k 个值是有序的，因此是适合求 topK 问题的数据结构。

由于存在重复值，因此采用 TreeMap 实现更合适，而不是 TreeSet。

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }

        TreeMap<Integer, Integer> map = new TreeMap<>();
        int cnt = 0;
        for (int num: arr) {
            if (cnt < k) {
                map.put(num, map.getOrDefault(num, 0) + 1);
                cnt++;
                continue;
            } 

            Map.Entry<Integer, Integer> entry = map.lastEntry();
            if (entry.getKey() > num) {
                map.put(num, map.getOrDefault(num, 0) + 1);
                if (entry.getValue() == 1) {
                    map.pollLastEntry();
                } else {
                    map.put(entry.getKey(), entry.getValue() - 1);
                }
            }
        }

        int[] res = new int[k];
        int idx = 0;
        for (Map.Entry<Integer, Integer> entry: map.entrySet()) {
            int freq = entry.getValue();
            while (freq-- > 0) {
                res[idx++] = entry.getKey();
            }
        }
        return res;
    }
}

4. 计数排序

当数据量比较小时，可以直接采用计数排序。

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0 || arr.length == 0) {
            return new int[0];
        }
        int[] counter = new int[10001];
        for (int num: arr) {
            counter[num]++;
        }
        int[] res = new int[k];
        int idx = 0;
        for (int num = 0; num < counter.length; num++) {
            while (counter[num]-- > 0 && idx < k) {
                res[idx++] = num;
            }
            if (idx == k) {
                break;
            }
        }
        return res;
    }
}