C++ TopK：使用 nth_element 快速求最小 K 个数（蓝桥杯 / LeetCode 可用）

题目链接

• 面试题 17.14. 最小 K 个数：面试题 17.14. 最小K个数 - 力扣（LeetCode）
• 215. 数组中的第 K 个最大元素（配套 TopK 题）：LCR 076. 数组中的第 K 个最大元素 - 力扣（LeetCode）

题目描述

设计一个算法，找出数组中最小的 k 个数。以任意顺序返回这 k 个数均可。

示例：

输入： arr = [1,3,5,7,2,4,6,8], k = 4
输出： [1,2,3,4]

提示：

• 0 <= arr.length <= 100000
• 0 <= k <= min(100000, arr.length)

剑指 Offer 40 和本题几乎一模一样，代码通用。

思路讲解

std::nth_element 是 <algorithm> 里的 STL 函数，不是全排序，只做分区 。 nth_element(first, pos, last) 保证：

1. *pos 的值就是数组排序后 pos 下标对应的值；
2. pos 左边所有元素 ≤ *pos，pos 右边所有元素 ≥ *pos；
3. 左右区间内部无序！

所以执行 nth_element(arr.begin(), arr.begin()+k-1, arr.end()) 之后： [0, k-1] 区间内就是整个数组最小的 k 个元素，直接截取前 k 个返回即可。

• 时间复杂度：平均 O (n)，是这道题理论上最快的写法
• 蓝桥杯、LeetCode GCC 均支持，可以放心使用

✅ 核心代码（nth_element 版本，推荐）

#include <bits/stdc++.h>
using namespace std;

class Solution {
public:
    vector<int> smallestK(vector<int>& arr, int k) {
        if (k == 0) return {};
        nth_element(arr.begin(), arr.begin() + k - 1, arr.end());
        return vector<int>(arr.begin(), arr.begin() + k);
    }
};

比赛推荐万能头 #include <bits/stdc++.h>，不用单独引入 vector /algorithm。

---

📌 备选方案

方案 1：partial_sort（返回的 k 个数有序）

如果题目要求输出的 k 个数字升序，就用这个，复杂度 \(O(n\log k)\)

#include <bits/stdc++.h>
using namespace std;

class Solution {
public:
    vector<int> smallestK(vector<int>& arr, int k) {
        if (k == 0) return {};
        partial_sort(arr.begin(), arr.begin() + k, arr.end());
        return vector<int>(arr.begin(), arr.begin() + k);
    }
};

方案 2：sort 全排序（最简单，大数据较慢）

\(O(n\log n)\)，数据量小的时候写起来最省事

#include <bits/stdc++.h>
using namespace std;

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

方案 3：大顶堆 priority_queue（适合数据流、超大数组）

\(O(n\log k)\)，不需要一次性加载全部数据时优先选择

#include <bits/stdc++.h>
using namespace std;

class Solution {
public:
    vector<int> smallestK(vector<int>& arr, int k) {
        if (k == 0) return {};
        // 大顶堆，维护k个最小值
        priority_queue<int> heap;
        for (int x : arr) {
            heap.push(x);
            if (heap.size() > k) {
                heap.pop();
            }
        }
        vector<int> res;
        while (!heap.empty()) {
            res.push_back(heap.top());
            heap.pop();
        }
        return res;
    }
};

⚠️ 重点注意事项

1. nth_element 不会排序前 k 个元素！ 判题允许无序的时候才能 AC；如果题目明确要求结果升序，必须换成partial_sort。

2. 原地修改原数组 如果业务不能破坏原始数组，先拷贝一份：

   vector<int> tmp = arr;
   nth_element(tmp.begin(), tmp.begin() + k - 1, tmp.end());

3. 边界 k=0 必须处理 不单独判断，arr.begin() -1 会越界，产生未定义行为。

4. 蓝桥杯可以正常使用 nth_element，属于 C++ 标准库，没有禁用。

5. 面试的时候，如果面试官不让调用 STL 内置函数，需要手写快速选择（QuickSelect）。

📊 四种方案对比

表格

方案	时间复杂度	前 k 个是否有序	适用场景
nth_element	平均 O (n)	❌ 无序	LeetCode、蓝桥杯，追求最快
partial_sort	O(n log k)	✅ 有序	需要输出有序结果
sort 全排序	O(n log n)	✅ 全部有序	数据量很小，图省事
大顶堆	O(n log k)	❌ 堆输出逆序	数据流、海量数据