题解 | 最小的K个数

1、解题思路

1. 排序法：直接对整个数组排序，然后取前 k 个元素。时间复杂度为 O(nlogn)，不满足题目要求的 O(nlogk)。
2. 堆（优先队列）：使用一个大根堆来维护当前最小的 k 个元素。遍历数组，将每个元素与堆顶（当前堆中的最大值）比较： 如果堆的大小小于 k，直接加入堆。如果当前元素小于堆顶，弹出堆顶并加入当前元素。最终堆中的元素即为最小的 k 个数。时间复杂度为 O(nlogk)，空间复杂度为 O(k)。
3. 快速选择（Quickselect）：类似于快速排序的分区思想，可以在平均 O(n) 时间内找到第 k 小的元素，然后输出前 k 个元素。最坏情况下时间复杂度为 O(n^2)，但题目要求 O(nlogk)，因此不推荐。

2、代码实现

C++

#include <queue>
#include <vector>
class Solution {
  public:
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     *
     * @param input int整型vector
     * @param k int整型
     * @return int整型vector
     */
    vector<int> GetLeastNumbers_Solution(vector<int>& input, int k) {
        // write code here
        vector<int> res;
        if (k <= 0 || k > input.size()) {
            return res;
        }

        priority_queue<int> maxHeap;    // 大根堆
        for (int num : input) {
            if (maxHeap.size() < k) {
                maxHeap.push(num);
            } else if (num < maxHeap.top()) {
                maxHeap.pop();
                maxHeap.push(num);
            }
        }

        while (!maxHeap.empty()) {
            res.push_back(maxHeap.top());
            maxHeap.pop();
        }

        return res;
    }
};

Java

import java.util.*;


public class Solution {
    /**
     * 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
     *
     *
     * @param input int整型一维数组
     * @param k int整型
     * @return int整型ArrayList
     */
    public ArrayList<Integer> GetLeastNumbers_Solution (int[] input, int k) {
        // write code here
        ArrayList<Integer> res = new ArrayList<>();
        if (k <= 0 || k > input.length) {
            return res;
        }

        PriorityQueue<Integer> maxHeap = new PriorityQueue<>(Collections.reverseOrder());
        for (int num : input) {
            if (maxHeap.size() < k) {
                maxHeap.offer(num);
            } else if (num < maxHeap.peek()) {
                maxHeap.poll();
                maxHeap.offer(num);
            }
        }

        while (!maxHeap.isEmpty()) {
            res.add(maxHeap.poll());
        }
        return res;
    }
}

Python

#
# 代码中的类名、方法名、参数名已经指定，请勿修改，直接返回方法规定的值即可
#
# 
# @param input int整型一维数组 
# @param k int整型 
# @return int整型一维数组
#
import heapq

class Solution:
    def GetLeastNumbers_Solution(self , input: List[int], k: int) -> List[int]:
        # write code here
        if k <= 0 or k > len(input):
            return []
        
        maxHeap = []
        for num in input:
            if len(maxHeap) < k:
                heapq.heappush(maxHeap, -num)  # 用负数模拟大根堆
            elif num < -maxHeap[0]:
                heapq.heappop(maxHeap)
                heapq.heappush(maxHeap, -num)
        
        res = [-x for x in maxHeap]
        return res

3、复杂度分析

1. 时间复杂度：O(nlogk)，因为每个元素最多插入和弹出堆一次，堆的操作时间为 O(logk)。
2. 空间复杂度：O(k)，堆的大小最多为 k。