最大的 K 个数 - TopK 问题

TopK 问题的经典方法有两种, 本文主要记录 快速选择的方法

采用堆的方法详见 TopK 问题 - 堆的方法

问题

从一个无序数组中,找出最大的 $k$ 个数。

例如:[5, 1, 2, 4, 3] 中最大的两个数是 [5, 4]

快速选择的方法

快速选择的方法,类似 快速排序 的思路:

  1. 对数组进行原地分割。

    选择一个基准元素,分割后使得:

    大于它的元素都在它左边,小于它的元素都在它右边。

  2. 根据分割后左边的元素个数和 $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$ 进行比较,有三种情况。

  1. 如果 $k = m$ ,那么已经找到最大的 $k$ 个数,终止递归。

  2. 如果 $k < m$ ,说明最大的 $k$ 个数一定在左边的 $m$ 个数之中。

    但是注意,左边的 $m$ 个数只保证不大于基准元素,并非排序。

    此时,需要对左边的 $m$ 个数递归分割。

  3. 如果 $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$ 。

它不同于快速排序的是, 快速排序每次都递归分割两边,而快速选择每次只选择一边递归。

(完)

* * *
评论 首页 订阅 打印 ↑顶部