题目要求「任意顺序返回这 k 个数即可」，可以利用快速排序的 partition 操作。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
 public:
  vector<int> smallestK(vector<int>& arr, int k) {
    if (k == 0) return {};
    int l = partition(arr, 0, arr.size() - 1, k);
    return vector<int>(arr.begin(), arr.begin() + l + 1);
  }
  int partition(vector<int>& a, int l, int r, int k) {
    if (l == r) return l;
    int mid = l + r >> 1;
    int p = a[mid];
    int i = l - 1, j = r + 1;
    while (i < j) {
      while (a[++i] < p);
      while (a[--j] > p);
      if (i < j) swap(a[i], a[j]);
    }
    int sl = j - l + 1;
    if (k <= sl)
      return partition(a, l, j, k);
    else
      return partition(a, j + 1, r, k - sl);
  }
};