面试时遇到的一道题:给出给出 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;
}

这道题在 力扣 上也有,看了题解之后发现这是算法导论上讲过的,感觉自己不仅算法水平低下,而且还读书不认真😭。