【优先队列, 排序, BST】最小的k个数

题目

剑指 Offer 40. 最小的k个数

---

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

解题

方法一：大根堆

思路

本题是经典的 Top K 问题(前 K 小)，维护一个容量为 k 的大根堆(maxHeap)。
先用数组中前 k 个元素把优先队列填满，接着遍历数组，如果当前元素小于堆顶，就把堆顶元素出队，然后放入当前元素，否则什么都不干。
遍历完成后把优先队列中所有元素取出放入数组中返回即可。

代码

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        if (k == 0) return new int[0];
        Queue<Integer> maxHeap = new PriorityQueue<>(Comparator.reverseOrder());
        for (int i = 0; i < k; ++i) maxHeap.offer(arr[i]);
        for (int i = k; i < arr.length; ++i) {
            if (arr[i] < maxHeap.peek()) {
                maxHeap.poll();
                maxHeap.offer(arr[i]);
            }
        }
        int[] ans = new int[k];
        for (int i = 0; i < k; ++i) ans[i] = maxHeap.poll();
        return ans;
    }
}

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

方法二：排序

思路

直接把数组升序排序，然后返回前 k 个元素即可。

代码

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        Arrays.sort(arr);
        return Arrays.copyOf(arr, k);
    }
}

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        sort(arr.begin(), arr.end());
        return {arr.begin(), arr.begin() + k};
    }
};

本题数据范围比较小，直接上计数排序。

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        int[] counts = new int[10001];
        for (int num : arr) ++counts[num];
        int[] topK = new int[k];
        for (int i = 0; i <= 10000 && k > 0; ++i) {
            while (counts[i]-- > 0 && k > 0) topK[--k] = i;
        }
        return topK;
    }
}

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        int counts[10001];
        memset(counts, 0, sizeof(counts));
        for (int& num : arr) ++counts[num];
        vector<int> top_k;
        for (int i = 0; i <= 10000 && k > 0; ++i) {
            while (counts[i]-- && k) {
                top_k.push_back(i);
                --k;
            }
        }
        return top_k;
    }
};

方法三：二叉搜索树(BST)

思路

思路与上面差不多，也是利用二叉搜索树进行排序。

C++ 直接用 multiset，Java 里没有可重复的有序集合，使用 TreeMap 代替。

代码

class Solution {
public:
    vector<int> getLeastNumbers(vector<int>& arr, int k) {
        multiset<int> bst(arr.begin(), arr.end());
        auto pos_k = bst.begin();
        advance(pos_k, k);
        return vector<int>(bst.begin(), pos_k);
    }
};

class Solution {
    public int[] getLeastNumbers(int[] arr, int k) {
        Map<Integer, Integer> bst = new TreeMap<Integer, Integer>();
        for (int num : arr) bst.put(num, bst.getOrDefault(num, 0) + 1);
        int[] topK = new int[k];
        outer: for (var entry : bst.entrySet()) {
            int num = entry.getKey(), cnt = entry.getValue();
            while (cnt-- > 0) {
                if (--k < 0) break outer;
                topK[k] = num;
            }
        }
        return topK;
    }
}