TopK 问题的经典方法有两种, 本文主要记录 快速选择的方法。
采用堆的方法详见 TopK 问题 - 堆的方法 。
问题 ¶
从一个无序数组中,找出最大的 $k$ 个数。
例如:[5, 1, 2, 4, 3]
中最大的两个数是 [5, 4]
。
快速选择的方法 ¶
快速选择的方法,类似 快速排序 的思路:
-
对数组进行原地分割。
选择一个基准元素,分割后使得:
大于它的元素都在它左边,小于它的元素都在它右边。
-
根据分割后左边的元素个数和 $k$ 的大小关系,选择一边递归进行分割。
原地分割
以数组 [5,3,1,9,7,2,5,8,4]
为例,选择第一个元素为基准元素。
原地分割后,大于基准值的放左边,小于基准值的放右边。

原地分割算法和 快速排序 中的原地分割完全一致。
详细可以参考 原地分割 - 三份 , 其时间复杂度是 $O(n)$ 。
TopK 问题 - 快速选择 - 原地分割 - C 语言实现
// 数组原地分割,取 v = a[start]
// >v 在左,=v 在中,<v在右
// 返回基准元素在整个数组中的位置
// 和快排完全一样
int Partition(int a[], int start, int end) {
int v = a[start];
int left = start;
int right = end;
int i = start;
while (i <= right) {
if (a[i] > v) {
Swap(a, i, left);
left++;
i++;
} else if (a[i] < v) {
Swap(a, i, right);
right--;
} else {
i++;
}
}
return i - 1;
}
快速选择的主过程
原地分割函数返回的是基准元素在整个数组中的位置 p
。
取 m = p + 1
,那么 m
表示的就是整个数组中基准元素左边的元素个数。

想要求的是整个数组中最大的 $k$ 个数,那么让它和 $m$ 进行比较,有三种情况。
-
如果 $k = m$ ,那么已经找到最大的 $k$ 个数,终止递归。
-
如果 $k < m$ ,说明最大的 $k$ 个数一定在左边的 $m$ 个数之中。
但是注意,左边的 $m$ 个数只保证不大于基准元素,并非排序。
此时,需要对左边的 $m$ 个数递归分割。
-
如果 $k > m$ ,说明已经找到了最大的 $m$ 个数,但是还有 $k-m$ 个数在右边。
此时,需要对右边的 $m$ 个数递归分割。
由此思路,即得出实现代码。
TopK 问题 - 快速选择法 - C 语言实现
void QuickSelect(int a[], int start, int end, int k) {
if (start >= end || k <= 0) return;
int p = Partition(a, start, end);
int m = p + 1; // 整个数组中在基准元素左边的元素个数
if (k < m)
QuickSelect(a, start, p - 1, k);
else if (k > m)
QuickSelect(a, p + 1, end, k); // 注意传 k 而非 k - m ,对齐 m 的意义
else
return;
}
// TopK 会原地把 k 大数放到 a 的前 k 个
void TopK(int a[], int n, int k) { return QuickSelect(a, 0, n - 1, k); }
复杂度分析
递归深度是 $\log{n}$ ,原地操作,因此空间复杂度是 $O(\log{n})$ 。
每次都只会选择一边进行递归分割,平均来看:
- 第一次选择 $n$ 个元素进行分割
- 第二次选择 $n/2$ 个元素进行分割
- 第二次选择 $n/4$ 个元素进行分割
- ….
每一次的原地分割时间复杂度是线性的 $O(n)$, 因此总的时间复杂度是:
\[O(n + \frac {n} {2} + \frac {n} {4} + \ldots) \\= O(n (1 + \frac {1} {2} + \frac {n} {4} + \ldots )) \\\leq O(2n) \\= O(n)\]其中,系数是 几何级数 ,收敛到 $2$ 。
它不同于快速排序的是, 快速排序每次都递归分割两边,而快速选择每次只选择一边递归。
(完)