使用Quick Select算法求数组第K大元素遇大K超时问题求助

我在做LeetCode的「数组中的第K个最大元素」题目时，采用Quick Select算法实现，但遇到超大数组且K值很大（如50000）时出现超时。尝试过随机打乱数组避免固定基准的最坏情况，但问题仍未解决——大数组搭配小K时正常，仅大K场景超时，想知道耗时原因及优化建议。

我的代码如下：

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """

        n = len(nums)
        random.shuffle(nums)

        return self.quickselect(nums, 0, n-1, n-k)


    def quickselect(self, nums, l, r, k):

        piv = self.partition(nums, l, r)

        if k < piv:
            return self.quickselect(nums, l, piv-1, k)
        elif k > piv:
            return self.quickselect(nums, piv+1, r, k)
        else:
            return nums[piv]


    def partition(self, nums, l, r):
        
        piv = r
        i = l

        for j in range(l, r):

            if nums[j] <= nums[piv]:
                nums[i], nums[j] = nums[j], nums[i]
                i += 1

        nums[i], nums[piv] = nums[piv], nums[i]
        return i

问题分析与优化建议 #

1. 缺失random模块导入

代码中使用了random.shuffle但未导入random模块，在LeetCode环境中会触发NameError，虽不是超时直接原因，但会导致代码无法正常运行，需先添加import random。

2. 单路Partition在重复元素场景下的低效

当前的单路Partition将小于等于基准的元素移到左侧，当数组存在大量重复元素时，每次Partition仅能将基准元素放到正确位置，其余重复元素仍留在左右区间，导致递归/迭代次数剧增，时间复杂度退化为O(n²)——这是大K场景超时的核心原因。

3. 递归栈开销

递归实现的Quick Select在处理超大数组时，会产生额外的栈帧开销，尤其是K很大时（对应n-k很小，递归深度可达O(logn)），累积的栈开销会拖慢执行速度。

优化后的代码实现 #

改用三路快排Partition（将数组分为小于、等于、大于基准的三部分）+ 迭代实现，能有效处理重复元素并避免递归栈开销：

import random

class Solution(object):
    def findKthLargest(self, nums, k):
        """
        :type nums: List[int]
        :type k: int
        :rtype: int
        """
        n = len(nums)
        random.shuffle(nums)
        target = n - k  # 转换为找第target小的元素（0-based）
        l, r = 0, n - 1
        
        while l <= r:
            lt, gt = self.partition_three_way(nums, l, r)
            if target < lt:
                r = lt - 1
            elif target >= gt:
                l = gt + 1
            else:
                # target落在等于基准的区间内，直接返回
                return nums[target]
        return -1

    def partition_three_way(self, nums, l, r):
        pivot = nums[r]
        lt = l  # 小于pivot的区域右边界（开区间）
        gt = r  # 大于pivot的区域左边界（开区间）
        i = l   # 当前遍历指针
        
        while i < gt:
            if nums[i] < pivot:
                nums[i], nums[lt] = nums[lt], nums[i]
                lt += 1
                i += 1
            elif nums[i] > pivot:
                gt -= 1
                nums[i], nums[gt] = nums[gt], nums[i]
            else:
                # 等于基准，直接跳过
                i += 1
        # 将基准元素移到等于区间的末尾
        nums[gt], nums[r] = nums[r], nums[gt]
        return lt, gt

额外优化思路 #

• 基准选择优化：若想进一步降低最坏情况概率，可采用「中位数的中位数」策略选择基准，不过实现复杂度较高，一般三路Partition+随机打乱已能应对绝大多数场景。
• 反向Partition适配大K：当K很大时（如K > n/2），可直接按降序进行Partition，目标索引设为k-1，避免转换为小索引的左半部分递归，进一步提升效率。

内容的提问来源于stack exchange，提问作者Justin Jin