数组TOP-K问题：求前K个最小元素的多种解法与C++实现

目录

引言

问题定义

方法一：排序法（最直观）

思路

方法二：堆（优先队列）法

思路

复杂度分析：

方法三：快速选择法（Quick Select）

方法四：计数排序法（数据范围有限时）

方法对比与选择建议

总结

引言

在算法面试和日常开发中，TOP-K问题是一类非常经典的问题。今天我们来探讨其中一个变体：给定一个无序数组，找出其中最小的K个元素。这个问题看似简单，实则蕴含着丰富的算法思想，我们将从暴力解法到最优解，逐一剖析并给出C++实现。

问题定义

输入：无序数组 arr，整数 K (1 ≤ K ≤ arr.length) 输出：数组中最小的K个元素（顺序不限） 示例：arr = [3,2,1,5,6,4], K = 2 → 输出 [1,2] 或 [2,1]

方法一：排序法（最直观）

思路

最简单的思路：将数组排序，然后取前K个元素。这种方法虽然时间复杂度较高，但代码简单，适合K接近N或对性能要求不高的场景。

复杂度分析

时间复杂度：O(N log N)，主要来自排序

空间复杂度：O(1)（原地排序）或 O(N)（使用额外空间）

C++代码实现：

#include <vector> #include <algorithm> using namespace std; vector<int> topKBySort(vector<int>& nums, int k) { if (k >= nums.size()) return nums; sort(nums.begin(), nums.end()); return vector<int>(nums.begin(), nums.begin() + k); }

方法二：堆（优先队列）法

思路

维护一个大小为K的最大堆：

1). 遍历数组元素

2). 当堆的大小小于K时，直接入堆

3). 当堆的大小等于K时，如果当前元素小于堆顶（堆中最大元素），则弹出堆顶，将当前元素入堆

遍历结束后，堆中保存的就是最小的K个元素。

为什么用最大堆而不是最小堆？因为我们需要快速访问当前K个候选元素中的最大值，以便判断新元素是否需要加入。

复杂度分析：

时间复杂度：O(N log K)，每个元素都可能进行堆操作

空间复杂度：O(K)，堆的存储空间

C++代码实现：

#include <vector> #include <queue> using namespace std; vector<int> topKByHeap(vector<int>& nums, int k) { if (k >= nums.size()) return nums; if (k <= 0) return {}; // 使用最大堆（优先队列默认是最大堆） priority_queue<int> maxHeap; for (int num : nums) { if (maxHeap.size() < k) { maxHeap.push(num); } else if (num < maxHeap.top()) { maxHeap.pop(); maxHeap.push(num); } } // 将堆中的元素取出 vector<int> result; while (!maxHeap.empty()) { result.push_back(maxHeap.top()); maxHeap.pop(); } return result; }

方法三：快速选择法（Quick Select）

思路

基于快速排序的分区思想：

1). 选择一个基准元素（pivot）

2). 将数组分区，使得左边元素都小于等于pivot，右边都大于pivot

3). 如果pivot的位置恰好是K-1，则左边0到K-1就是前K小元素

4). 如果pivot的位置大于K-1，则在左边继续查找

5). 如果pivot的位置小于K-1，则在右边继续查找

这种方法能在平均情况下达到线性时间复杂度。

复杂度分析

时间复杂度：平均O(N)，最坏O(N²)（可通过随机化pivot优化）

空间复杂度：O(1)（原地操作，不考虑递归栈）

C++代码实现：

#include <vector> #include <cstdlib> #include <ctime> using namespace std; class QuickSelect { private: int partition(vector<int>& nums, int left, int right) { // 随机选择pivot，避免最坏情况 int randomIndex = left + rand() % (right - left + 1); swap(nums[randomIndex], nums[right]); int pivot = nums[right]; int i = left; // i指向小于等于pivot的区域边界 for (int j = left; j < right; j++) { if (nums[j] <= pivot) { swap(nums[i], nums[j]); i++; } } swap(nums[i], nums[right]); return i; // 返回pivot的最终位置 } void quickSelect(vector<int>& nums, int left, int right, int k) { if (left >= right) return; int pivotIndex = partition(nums, left, right); if (pivotIndex == k) { return; } else if (pivotIndex > k) { quickSelect(nums, left, pivotIndex - 1, k); } else { quickSelect(nums, pivotIndex + 1, right, k); } } public: vector<int> topKByQuickSelect(vector<int> nums, int k) { if (k >= nums.size()) return nums; if (k <= 0) return {}; srand(time(nullptr)); quickSelect(nums, 0, nums.size() - 1, k - 1); // 前k个元素就是最小的k个（但未排序） return vector<int>(nums.begin(), nums.begin() + k); } };

方法四：计数排序法（数据范围有限时）

思路

当数组元素的范围较小且已知时，可以使用计数排序的思想：

1). 统计每个元素出现的频率

2). 从小到大遍历可能的元素值，累加计数直到达到K个

这种方法在某些特定场景下效率极高。

复杂度分析

时间复杂度：O(N + range)，range为元素取值范围大小

空间复杂度：O(range)

C++代码实现：

#include <vector> #include <algorithm> using namespace std; vector<int> topKByCounting(vector<int>& nums, int k) { if (k >= nums.size()) return nums; if (k <= 0) return {}; // 假设元素范围在[0, 10000]之间 int maxVal = *max_element(nums.begin(), nums.end()); int minVal = *min_element(nums.begin(), nums.end()); vector<int> count(maxVal - minVal + 1, 0); // 统计频率 for (int num : nums) { count[num - minVal]++; } // 收集结果 vector<int> result; for (int i = 0; i < count.size(); i++) { while (count[i] > 0 && result.size() < k) { result.push_back(i + minVal); count[i]--; } if (result.size() == k) break; } return result; }

方法对比与选择建议

总结

TOP-K问题虽然基础，但涵盖了排序、堆、分治等多种算法思想。在实际应用中，我们应该根据数据规模和特点选择合适的方法：

• 如果K接近N，排序法简单直接
• 如果数据量巨大且K较小，堆方法是工业界最常用的方案
• 如果追求极致性能且能接受最坏情况，快速选择是不错的选择
• 如果数据范围有限，计数排序可能带来惊喜

理解这些方法不仅能解决TOP-K问题，更能帮助我们建立算法思维，应对更复杂的场景。