Top-k 问题

发表于 分类于 ， Valine：

Top k 问题是一个经典的算法问题，要求求解数组中第 k 大（or 第k小的元素）。最简单的，该问题可以通过排序来解决，算法的时间复杂度为 O(nlogn)。有没有效率更高的算法呢？

实际上，我们可以使用类似快排的思想，来降低时间复杂度。这里，我们使用 leetcode 中的题目作为例子（便于测试）215. 数组中的第K个最大元素。具体做法是：

  1. 在数组区间 [0,len) 中随机选择一个 pivot 值，然后遍历数组，如果数值大于 pivot 的值，则将这个数放置到数组的前半部分，否则放置到数组的后半部分。
  2. 完成第一步后，我们得到有 x 个元素大于 pivot 值，与 k 进行比较，如下三种情况：
    1. 情况1：如果 x == k-1, 则说明 pivot 为第 k 个最大的元素，直接返回 pivot。
    2. 情况2：如果 x < k-1, 则说明 第k大的元素大于等于 pivot 值，调整区间 左边界 left 等于 pivot 的下标加1，右边界不动；
    3. 情况3：如果 x > k-1, 则说明 第k大的元素小于 pivot 值，调整区间的右边界 right 等于 pivot 的下标。
    4. 重复循环，分析 [left, right) 区间的情况。

实现代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Solution {
    public int findKthLargest(int[] nums, int k) {
        Random rand = new Random();
        int left = 0, right = nums.length;
        int pivot, cur = -1;
        while(cur != k-1){
            // 将大于 nums[pivot] 的元素放置在 左边，小于 nums[pivot] 的元素放置在右边
            cur = left-1;
            pivot = left + rand.nextInt(right-left);
            // 先将选定的 枢值 放置到区间的最后
            swap(nums, pivot, right-1);
            int val = nums[right-1];
            // 将大于枢值的值调到区间的最前面，cur 记录大于 枢值部分的最后一个元素
            for(int i = left; i < right; ++i){
                if(nums[i] > val){
                    cur++;
                    swap(nums, cur, i);
                }
            }
            // 将枢值调到 cur 后面的一个位置，从而枢值前面的数字都大于枢值，后面的都小于等于 枢值
            cur++;
            swap(nums, cur, right-1);
            if(cur < k){
                left = cur + 1;
            }else{
                right = cur;
            }
        }
        return nums[cur];
    }

    void swap(int[] nums, int x, int y){
        int tmp = nums[x];
        nums[x] = nums[y];
        nums[y] = tmp;
    }
}

上面的代码中，需要考虑一个特殊的情况，即如果 kth 元素在一个重复数字的区间段内，是否能够正确的返回。通过分析，发现是可以的，这种特殊的情况与 分析中的情况2 是一致，即遍历到最后时，会有大于 枢值的元素个数为 0，即 x == 0，这时候，枢轴元素会被移动到最左边，即 left 指向的位置。最后，左区间会基于枢轴位置向右调整1格，即到 left + 1。因此，不会进入死循环，cur 会不断移动，直到等于 k-1 为止。

对于上述的算法，它的实际时间复杂度是与 pivot 值的选择有关的，在最坏的情况下，时间复杂度可能达到 O(n^2)。但通过 随机选择枢轴，从概率的角度来说，该算法的平均时间复杂度为 O(n)。具体的数学证明在这里省略了，我们可以通过 主定理 粗略的进行分析（因为主定理要求子问题规模相同，而这里并不满足，所以是说是粗略分析）：

  • 子问题的个数a 为 1个，因为我们将区间划分为两块，而最后只需要分析其中一块即可；
  • 输入数据规模虽小的因子b 为2，假设枢轴恰好位于数组中间；
  • 除递归操作额外所需要的操作耗时为 O(n), 即 d = 1。

因此 $a<b^d$, 算法的时间复杂度为 O(n^d) = O(n)。

类似这个类型的题目还有 347. 前 K 个高频元素。对于寻找前 k 个高频元素，首先需要将各个元素出现的次数统计到一个以元素值为 key，出现次数为 value 的 map 中。然后通过上述 方法查到前 k 个元素（前k个元素之间不需要按顺序排列！）。

一些阅读材料：

  1. System Design Interview - Top K Problem (Heavy Hitters)