Kth Largest Element in an Array

import heapq
import random


# Approach 1: Sort — O(n log n)
def find_kth_largest_sort(nums: list[int], k: int) -> int:
    return sorted(nums, reverse=True)[k - 1]


# Approach 2: Min-heap of size k — O(n log k)
def find_kth_largest_heap(nums: list[int], k: int) -> int:
    heap = nums[:k]
    heapq.heapify(heap)
    for n in nums[k:]:
        if n > heap[0]:
            heapq.heapreplace(heap, n)  # pop smallest, push n
    return heap[0]


# Approach 3: Quickselect — O(n) average
def find_kth_largest(nums: list[int], k: int) -> int:
    target = len(nums) - k  # k-th largest = (n-k)-th smallest

    def quickselect(lo, hi):
        pivot = nums[random.randint(lo, hi)]  # random pivot avoids O(n²) worst case
        lt, eq, gt = lo, lo, hi  # three-way partition (handles duplicates)

        while eq <= gt:
            if nums[eq] < pivot:
                nums[lt], nums[eq] = nums[eq], nums[lt]
                lt += 1
                eq += 1
            elif nums[eq] > pivot:
                nums[eq], nums[gt] = nums[gt], nums[eq]
                gt -= 1
            else:
                eq += 1

        # pivot occupies indices [lt..gt]
        if target < lt:
            return quickselect(lo, lt - 1)
        elif target > gt:
            return quickselect(gt + 1, hi)
        else:
            return nums[target]

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


if __name__ == '__main__':
    nums = [3, 2, 1, 5, 6, 4]
    print(find_kth_largest_sort(nums[:], 2))  # 5
    print(find_kth_largest_heap(nums[:], 2))  # 5
    print(find_kth_largest(nums[:], 2))       # 5