第k个最大的数

摘抄自 leetcode LCR 076

这个问题我在《算法分析与设计》课上听过，今天偶然看到有人在 leetcode 上做了这道题，就想写篇文章梳理下思路。

问题简介

给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。

请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。

思路分析

首先最容易想到的方法就是排序之后直接选位于第k的数，但是如果真是这样这道题就不会是中等难度。

这种方法的时间复杂度是O(nlogn)，要超过这种方法，意思是我们的目标是O(n)的复杂度。

接下来说明解决思路

快速选择算法（分治）

在课上，这道题被作为一个经典的使用分治法的例子。

让我们用分治的思想来试一下，看如何拆分这个问题。

对于一个数组做分治，自然会想到把一个数组分成两个子数组，那么问题的关键就是如何划分。

对于这道题，我们便需要按照数组元素大小进行划分。

我们可以先在数组中取一个值 value ，将数组划分未小于 value 的 small 数组、等于 value 的 equal数组 和大于 value 的 large 数组。

然后统计 large 数组的元素数量 len_large 和 equal 数组的元素数量 len_equal。

如果 len_large $\ge$ k，说明我们要找的数在 large 数组中，那么我们只需要在 large 数组中继续进行分治，递归地找第 k 大的数。

如果 len_large $<$ k $\le$ len_large+len_equal，说明我们要找的数在 equal 数组中，说明 value 就是我们要找的数，直接返回结果。

最后，其它的情况，说明我们要找的数在 small 数组中，那么我们只需要在 small 数组中继续进行分治，递归地找第 k-len_large-len_equal 大的数。

代码实现

python复制
from line_profiler import profile
import random

先写个使用排序的版本的（为了排除语言导致的算法速度差异，使用python语言实现的快速排序）

python复制
def partition(arr, low, high):
    i = (low - 1) # 最小元素索引
    pivot = arr[high] # 基准值

    for j in range(low, high):
        # 当前元素小于或等于 pivot
        if arr[j] <= pivot:
            i = i + 1
            arr[i], arr[j] = arr[j], arr[i]

    arr[i + 1], arr[high] = arr[high], arr[i + 1]
    return (i + 1)


# 快速排序函数
def quickSort(arr, low, high):
    if low < high:
        pi = partition(arr, low, high)
        quickSort(arr, low, pi - 1)
        quickSort(arr, pi + 1, high)


def quick_sort(nums):
    quickSort(nums, 0, len(nums) - 1)
    return nums


sorted = quick_sort # 覆盖掉标准库的排序函数，排除语言速度的干扰


@profile
def sort_find_kth_largest(nums, k):
    sorted(nums)
    return nums[-k]

python复制
@profile
def find_kth_largest(nums, k):
    value = nums[0]
    large = [x for x in nums if x > value]
    equal = [x for x in nums if x == value]
    small = [x for x in nums if x < value]

    if k <= len(large):
        return find_kth_largest(large, k)
    elif k > len(large) + len(equal):
        return find_kth_largest(small, k - len(large) - len(equal))
    else:
        return value

python复制
def test():
    l = 1000
    nums = [random.randint(1, l) for _ in range(l)]
    k = random.randint(1, l)
  
    sort_find_kth_largest(nums, k) # 测试包好像有点bug，先算一次，消除这个影响

    print(find_kth_largest(nums, k))
    print(sort_find_kth_largest(nums, k))

%load_ext line_profiler
%lprun -f test test()

plaintext复制
    612  
    612  
    Timer unit: 1e-07 s    Total time: 1.59446 s
    Function: test at line 1    Line #      Hits         Time  Per Hit   % Time  Line Contents
    ==============================================================
         1                                           def test():
         2         1         10.0     10.0      0.0      l = 1000
         3         1     108511.0 108511.0      0.7      nums = [random.randint(1, l) for _ in range(l)]
         4         1        117.0    117.0      0.0      k = random.randint(1, l)
         5   
         6         1     233049.0 233049.0      1.5      sort_find_kth_largest(nums, k) # 测试包好像有点bug，先算一次，消除这个影响
         7   
         8         1    4716746.0    5e+06     29.6      print(find_kth_largest(nums, k))
         9         1   10886168.0    1e+07     68.3      print(sort_find_kth_largest(nums, k))

可以看到，使用分治法的快速选择算法还是快很多的。

但是这种算法还是有缺陷的，我们通过时间复杂度分析来揭示一下这个缺陷。

时间复杂度分析

这是一个使用递归的解法，设时间复杂度为 $T(n)$，每次划分数组，较大部分的比例为 $p, p<1$，则有：

$T(n) = n + p·T(pn) + (1-p)·T((1-p)n)$

当 $p = \frac{1}{2}$ 时

使用 Master Theorem 计算 $T(n)$ 得时间复杂度为 $O(n)$

但是当 $p = 1 或 0$ 时

时间复杂度为 $O(nlogn)$

所以会导致实际使用时时间复杂度和 $O(nlogn)$ 差不太多，并不理想。(下面是leetcode给的时间复杂度分析图)

（顺便吐槽下leetcode用python写题怎么可以用标准库的sort，速度直接完爆了我用python原始实现的快速选择算法）

为了避免这种情况，引出下面的中位数的中位数算法

中位数的中位数算法

这个算法的目的是为了以一个以 O(1) 的时间复杂度找到数组中接近中位数的数。

算法流程如下

1. 将输入数组的元素划分为组，每组5个元素，且至多只有一个组由剩下的元素组成。
2. 通过将每个组进行快速排序，然后获取这至多由5个元素组成的组的中位数
3. 对于 2 中获取到的中位数，递归进行 1 和 2，直到最后只得到一个数，即为最终找到的中位数。
4. 将这个中位数作为基准，使用上面提到的快速选择算法，返回最终的结果。

代码实现

python复制
@profile
def median_find_kth_largest(nums, k):
    if len(nums) <= 5:
        return sorted(nums)[-k] # 这里使用的是被覆盖掉的排序算法，并非标准库
  
    # 划分数组为最多5个元素的子数组
    medians = []
    for i in range(0, len(nums), 5):
        group = nums[i:i+5]
        group_sorted = sorted(group)
        median = group_sorted[len(group)//2]  # 取中位数
        medians.append(median)

    # 这里改进为将中位数的中位数作为选择的value
    value = median_find_kth_largest(medians, len(medians)//2)

    large = [x for x in nums if x > value]
    equal = [x for x in nums if x == value]
    small = [x for x in nums if x < value]

    if k <= len(large):
        return median_find_kth_largest(large, k) # 全部调用自己，使用改进后的算法
    elif k > len(large) + len(equal):
        return median_find_kth_largest(small, k - len(large) - len(equal))
    else:
        return value

python复制
def new_test():
    l = 1000
    nums = [random.randint(1, l) for _ in range(l)]
    k = random.randint(1, l)
  
    sort_find_kth_largest(nums, k) # 测试包好像有点bug，先算一次，消除这个影响

    print(median_find_kth_largest(nums, k))
    print(find_kth_largest(nums, k))
    print(sort_find_kth_largest(nums, k))

%reload_ext line_profiler
%lprun -f new_test new_test()

plaintext复制
    408
    408
    408  
    Timer unit: 1e-07 s    Total time: 0.469127 s
    Function: new_test at line 1    Line #      Hits         Time  Per Hit   % Time  Line Contents
    ==============================================================
         1                                           def new_test():
         2         1          6.0      6.0      0.0      l = 1000
         3         1      38925.0  38925.0      0.8      nums = [random.randint(1, l) for _ in range(l)]
         4         1         41.0     41.0      0.0      k = random.randint(1, l)
         5   
         6         1      71064.0  71064.0      1.5      sort_find_kth_largest(nums, k) # 测试包好像有点bug，先算一次，消除这个影响
         7   
         8         1     147002.0 147002.0      3.1      print(median_find_kth_largest(nums, k))
         9         1    1216938.0    1e+06     25.9      print(find_kth_largest(nums, k))
        10         1    3217295.0    3e+06     68.6      print(sort_find_kth_largest(nums, k))

压倒性的胜利，可以看出这个算法已经和之前的算法不是一个数量级了。

最后我们再来分析一下这个算法的时间复杂度

时间复杂度分析（由DeepSeek讲解，个人已消化，自觉无法讲得更好）

中位数的中位数（Median of Medians）算法的时间复杂度为严格的 O(n)，即使是最坏情况。以下是详细分析：

---

递归过程的时间复杂度推导

设输入规模为 n，算法的时间复杂度为 T(n)。根据算法逻辑，递归过程分为以下步骤：

1. 分组并计算中位数：

   • 将数组划分为 ⌈n/5⌉ 组，每组最多5个元素。
   • 对每组排序（时间复杂度为 O(1)，因为每组最多5个元素）。
   • 收集每组的中位数，得到规模为 n/5 的中位数数组。
   • 时间复杂度：O(n)（每组排序的总时间为 O(n)）。

2. 递归找基准值（pivot）：

   • 对中位数数组递归调用 median_of_medians，找到其中位数作为基准值。
   • 子问题规模：T(n/5)。

3. 划分数组：

   • 根据基准值将原数组分为 left（< pivot）、mid（= pivot）、right（> pivot）三部分。
   • 时间复杂度：O(n)（遍历一次数组）。

4. 递归查找目标子数组：

   • 根据 k 的位置，选择在 left 或 right 中继续查找。
   • 由算法性质保证，无论 k 在左侧还是右侧，子问题规模最多为 7n/10（证明见下文）。
   • 子问题规模：T(7n/10)。

综上，递归式为：

$$T(n) = T\left(\frac{n}{5}\right) + T\left(\frac{7n}{10}\right) + O(n)$$

---

严格证明时间复杂度为 O(n)

通过数学归纳法或递归树分析可证明递归式满足 T(n) = O(n)。
核心观察：每次递归的总子问题规模递减，且总工作量为线性叠加。

关键不等式

$$T(n) \leq c \cdot n \quad \text{（假设）}$$

代入递归式：

$$c \cdot n \leq c \cdot \frac{n}{5} + c \cdot \frac{7n}{10} + k \cdot n$$

化简得：

$$c \cdot n \leq c \cdot \left( \frac{n}{5} + \frac{7n}{10} \right) + k \cdot n = c \cdot \frac{9n}{10} + k \cdot n$$

移项后：

$$c \cdot n - c \cdot \frac{9n}{10} \leq k \cdot n \implies c \cdot \frac{n}{10} \leq k \cdot n \implies c \leq 10k$$

因此，存在常数 c（如 c = 10k）使得不等式成立，即 T(n) = O(n)。

---

为何子问题规模最多为 7n/10？

• 中位数的中位数（pivot）保证至少 30% 的元素被淘汰：
  • 每组5个元素的中位数中，至少有3个元素 ≥ 中位数（或 ≤ 中位数）。
  • 因此，中位数的中位数（pivot）在整个数组中至少大于等于 3n/5 * 1/2 = 3n/10 个元素，也至少小于等于 3n/10 个元素。
  • 最坏情况下，子问题规模最多为 n - 3n/10 = 7n/10。

---

空间复杂度

• 递归栈深度：O(log n)（每次递归规模减少为 7n/10）。
• 临时数组：O(n)（存储 left、mid、right）。
• 总空间复杂度：O(n)。

---

总结

• 时间复杂度：O(n)（严格线性，无最坏情况退化）。
• 优势：相比快速选择（Quickselect）的 O(nlogn) 最坏时间复杂度，中位数的中位数算法通过精心选择基准值，保证了线性时间复杂度。