Learning On the Way

215. Kth Largest Element in an Array

Given an integer array nums and an integer k, return the kth largest element in the array. Note that it is the kth largest element in the sorted order, not the kth distinct element.

Quick select

class Solution:
    def findKthLargest(self, nums: List[int], k: int) -> int:          
        def partition(start, end):
            # Add some randomization
            idx = randint(start, end)
            nums[idx], nums[start] = nums[start], nums[idx]

            pivot = start
            pivot_value = nums[start]
            for i in range(start + 1, end + 1): # It is the same whether you start from "start" or "start + 1"
                if nums[i] < pivot_value:
                    pivot += 1
                    nums[i], nums[pivot] = nums[pivot], nums[i]

            nums[start], nums[pivot] = nums[pivot], nums[start]
            return pivot # Pivot is also the same as the number of elements smaller than the pivot element
        
        def kthSmallest(start, end, k):  # k start from 1
            pivot = partition(start, end)
            if pivot - start == k - 1:
                return nums[pivot]
            elif pivot - start > k - 1:
                return kthSmallest(start, pivot - 1, k)
            else:
                return kthSmallest(pivot + 1, end, k - pivot - 1 + start)

        return kthSmallest(0, len(nums) - 1, len(nums) + 1 - k)