Kth smallest元素大根堆Python解法Time limit exceeded问题求助

我在练习GeeksForGeeks平台的Kth smallest element题目，题目详情如下：

给定数组 arr[] 和整数 K，其中K小于数组长度，任务是查找给定数组中的**第K小（Kth smallest）**元素，题目保证所有数组元素互不相同。

预期时间复杂度（Expected Time Complexity）: O(n)
预期辅助空间（Expected Auxiliary Space）: O(log(n))
约束条件（Constraints）:

• 1 <= N <= 10^5
• 1 <= arr[i] <= 10^5
• 1 <= K <= N

我的实现代码 #

class Solution:
    def kthSmallest(self,arr, l, r, k):
        '''
        arr : given array
        l : starting index of the array i.e 0
        r : ending index of the array i.e size-1
        k : find kth smallest element and return using this function
        '''
        arr2=arr[:k]
        arr2.insert(0,None)
        for i in range(k//2,0,-1):
            arr2=self.heapify(arr2,i,k-1)
            
        for i in arr[k:]:
            if i <arr2[1]:
                arr2[1]=i
                arr2=self.heapify(arr2,1,k-1)
        return arr2[1]


    def heapify(self,arr, i, r):
        if 2 * i <= r + 1 and arr[2 * i] > arr[i]:
            arr[2 * i], arr[i] = arr[i], arr[i * 2]
            arr = self.heapify(arr, 2 * i, r)
        if 2 * i + 1 <= r + 1 and arr[2 * i + 1] > arr[i]:
            arr[2 * i + 1], arr[i] = arr[i], arr[i * 2 + 1]
            arr = self.heapify(arr, 2 * i + 1, r)
        return arr

我的实现思路为：首先截取数组前K个元素构建子数组，对其执行大根堆化（max heapify）操作；随后遍历数组剩余元素，若当前元素小于堆顶元素，则替换堆顶元素后对堆顶执行max heapify操作。运行代码时出现**时间超限（Time limit exceeded）**错误，需要定位问题诱因与优化方案。

问题原因与优化方案 #

超时核心原因 #

• 递归实现堆化的开销过高：Python函数递归调用的栈帧开销远大于迭代逻辑，在1e5量级的数据规模下，频繁递归会大幅拖慢运行速度。
• 存在额外的线性时间开销：arr2.insert(0, None)是数组头部插入操作，时间复杂度为O(k)，当k接近1e5时这一步就会消耗大量运行时间；每次递归返回后重新给arr2赋值也会产生不必要的对象操作开销。
• 堆化逻辑有冗余：左子节点交换完成后直接递归下沉，之后没有更新当前节点的最大值状态就直接判断右子节点，会产生重复比较，增加无效计算。
• 算法复杂度不满足题目要求：大小为K的大根堆找第K小的时间复杂度是O(n logk)，当K和数组长度同量级时，复杂度接近O(n logn)，达不到题目要求的O(n)时间复杂度，自然容易触发时间限制。

可落地的优化方案 #

1. 堆逻辑优化（可通过用例，但不满足最优复杂度要求）

去掉头部占位的冗余插入操作，把递归堆化改成迭代实现，堆化时每次统一比较当前节点和左右子节点的最大值，只做一次交换后循环下沉，消除递归和冗余判断的开销。
优化后的堆实现代码：

class Solution:
    def kthSmallest(self,arr, l, r, k):
        def max_heapify(heap, i, heap_size):
            while True:
                left = 2 * i + 1
                right = 2 * i + 2
                largest = i
                if left < heap_size and heap[left] > heap[largest]:
                    largest = left
                if right < heap_size and heap[right] > heap[largest]:
                    largest = right
                if largest == i:
                    break
                heap[i], heap[largest] = heap[largest], heap[i]
                i = largest
        
        heap = arr[:k]
        # 原地建大根堆
        for i in range(k//2 - 1, -1, -1):
            max_heapify(heap, i, k)
        # 遍历剩余元素维护堆
        for num in arr[k:]:
            if num < heap[0]:
                heap[0] = num
                max_heapify(heap, 0, k)
        return heap[0]

2. 快速选择算法（完全匹配题目复杂度要求）

基于快速排序分区思想实现快速选择，每次随机选取基准值将数组划分为小于基准、大于基准两部分，根据基准最终落点和K的位置关系，只需要递归处理其中一个分区，平均时间复杂度O(n)，递归栈空间平均O(logn)，完全符合题目要求。随机选基准的操作可以避免有序数组下的最坏时间复杂度问题。
快速选择实现代码：

import random
class Solution:
    def kthSmallest(self,arr, l, r, k):
        target = k - 1
        while l <= r:
            # 随机选基准规避最坏情况
            pivot_pos = random.randint(l, r)
            arr[pivot_pos], arr[r] = arr[r], arr[pivot_pos]
            pivot = arr[r]
            i = l
            # 分区：小于基准的放左边，大于的放右边
            for j in range(l, r):
                if arr[j] < pivot:
                    arr[i], arr[j] = arr[j], arr[i]
                    i += 1
            arr[i], arr[r] = arr[r], arr[i]
            # 匹配目标位置直接返回
            if i == target:
                return arr[i]
            elif i < target:
                l = i + 1
            else:
                r = i - 1

内容的提问来源于stack exchange，提问作者Debjyoti Sutradhar