面试时遇到的一道题:给出给出 n 个数,求出前 k 小的数字。
输入是 n + 1 行:
- 第一行是 n 和 k
- 之后 n 行是这个数列
输出是 k 个数。
当时没想出来怎么做,直接嗯排序,然后输出,OJ 的时限比较宽容,竟然过了😄。
不过在面试官那里肯定是没过的。
直接排序的时间复杂度是 $O(nlg{n})$ ,如果维护一个大小为 k 的堆则是能够优化成 $O(nlg{k})$ ,肯定有更好的方法。
后来想到了类似快排的划分的思想,代码和思路如下:
#include <iostream>
#include <vector>
using namespace std;
// 打印范围在 [begin, end) 的数组切片
void print_slice(vector<int> &a, int begin, int end) {
// cout << "-------" << endl;
for (int i = begin; i < end; ++i) {
cout << a[i] << " ";
}
// cout << endl;
// cout << "-------" << endl;
}
// 类似二分的思想,对于范围在 [begin, end) 的数组切片进行划分
void sol(vector<int> &a, int begin, int end, int k) {
if (begin >= end || k <= 0) {
return;
}
int lower = begin;
int upper = end - 1;
// 以最左边的为轴
// 小于它的放左边,大于它的放右边
int pivot = a[lower];
while (lower < upper) {
while (lower < upper && a[upper] > pivot) {
upper -= 1;
}
if (lower < upper) {
a[lower] = a[upper];
lower += 1;
}
while (lower < upper && a[lower] < pivot) {
lower += 1;
}
if (lower < upper) {
a[upper] = a[lower];
upper -= 1;
}
}
a[lower] = pivot;
int left_count = lower - begin;
// 左边的数就是前 left_count 小的数
// 如果轴左边数的数量刚好为 k - 1 ,则说明已经找到了前 k 小的数
// 如果小于 k - 1 ,则说明还需要在右边找前 k - 1 - left_count 小的数
// 如果大于 k - 1 ,则说明需要在左边找前 k 小的数
if (left_count <= k - 1) {
print_slice(a, begin, lower + 1);
sol(a, lower + 1, end, k - left_count - 1);
} else {
sol(a, begin, lower, k);
}
}
int main() {
int n, k;
cin >> n >> k;
vector<int> a(n);
for (int i = 0; i < n; ++i) {
cin >> a[i];
}
sol(a, 0, n, k);
cout << endl;
return 0;
}
这道题在 力扣 上也有,看了题解之后发现这是算法导论上讲过的,感觉自己不仅算法水平低下,而且还读书不认真😭。