leetcode-215 Kth Largest Element in an Array

# Given an integer array nums and an integer k, return the kth largest element.
# Note that it is the kth largest element in sorted order, not the kth distinct.
# You must solve it in O(n) time complexity.
#
# Example 1:
# Input: nums = [3,2,1,5,6,4], k = 2
# Output: 5
#
# Example 2:
# Input: nums = [3,2,3,1,2,4,5,5,6], k = 4
# Output: 4
#
# Constraints:
# 1 <= k <= nums.length <= 10^5
# -10^4 <= nums[i] <= 10^4

import heapq

class Solution:
    # Approach 1: Min-Heap of size k
    def findKthLargest(self, nums: List[int], k: int) -> int:
        # Maintain a min-heap of size k
        # The smallest element in this heap is the kth largest overall
        heap = []

        for num in nums:
            heapq.heappush(heap, num)
            # Keep heap size at k by removing smallest
            if len(heap) > k:
                heapq.heappop(heap)

        # Root of min-heap is kth largest
        return heap[0]

    # Approach 2: Using heapq.nlargest (cleaner but less educational)
    def findKthLargest2(self, nums: List[int], k: int) -> int:
        return heapq.nlargest(k, nums)[-1]

    # Approach 3: QuickSelect - O(n) average time
    def findKthLargest3(self, nums: List[int], k: int) -> int:
        # Convert to kth smallest problem
        k = len(nums) - k

        def quickselect(left, right):
            pivot = nums[right]
            p = left

            # Partition: elements < pivot go to left
            for i in range(left, right):
                if nums[i] <= pivot:
                    nums[p], nums[i] = nums[i], nums[p]
                    p += 1

            nums[p], nums[right] = nums[right], nums[p]

            if p == k:
                return nums[p]
            elif p < k:
                return quickselect(p + 1, right)
            else:
                return quickselect(left, p - 1)

        return quickselect(0, len(nums) - 1)

# Heap: Time O(n log k), Space O(k)
# QuickSelect: Time O(n) average, O(n^2) worst; Space O(1)
#
# Key Pattern: For "kth largest/smallest" problems, think heap!
# Min-heap of size k for kth largest
# Max-heap of size k for kth smallest