Leetcode215. Kth Largest Element in an Array-快速选择算法

题目

题目链接

思路

快速排序算法每次可以将一个元素放置到排序后的正确位置,因此每次交换后判断这个位置是否是第K大即可

复杂度

时间复杂度

• 第一次运行，pivot与n个元素比较，将输入规模减半并递归$T(n)=T(\frac{n}{2})+n$
• 第二次运行，pivot与$\frac{n}{2}$个元素比较，将输入规模减半并递归$T(\frac{n}{2})=T(\frac{n}{4})+\frac{n}{2}$
• $T(\frac{n}{4})=T(\frac{n}{8})+\frac{n}{4}$
  …
• $T(2)=T(1)+2$
• $T(1)=1$
• 则$T(n)=1+2+...+\frac{n}{8}+\frac{n}{4}+\frac{n}{2}+n=2\ast n$，$O(T(n))=O(n)$

空间复杂度

不考虑递归的栈开销，则没有空间消耗，$O(1)$

代码

class Solution:
    # quick select, based on quick sort
    def findKthLargest(self, nums: List[int], k: int) -> int:
        return self.quick_select(nums, 0, len(nums) - 1, k)
    
    def quick_select(self, nums, low, high, k):
        pivot = low
        
        for j in range(low, high):
            if nums[j] <= nums[high]:
                nums[pivot], nums[j] = nums[j], nums[pivot]
                pivot += 1
            
        nums[pivot], nums[high] = nums[high], nums[pivot]
        count = high - pivot + 1
        if count == k:  # 判断是否是第K大
            return nums[pivot]
        
        if count > k:
            return self.quick_select(nums, pivot + 1, high, k)
        return self.quick_select(nums, low, pivot - 1, k - count)