Find kth largest element

Solution 1

This solution uses a heap of size k.

1) Initialise a Heap Map

2) Add the k first values in the heap.

3) For the next element, check their value vs the smallest of the heap, if greater pop the smallest value and push the current one

4) Return the smallest value of the heap

heapq push and pop in log(n) time complexity as the undelying structure is a tree

import heapq

def findKthLargest(nums: List[int], k: int) -> int:
    heap = []
    heapq.heapify(heap)

    for i, num in enumerate(nums):
        # while less than k elements are in the heap just add new elements
        if i < k:
            heapq.heappush(heap,num)
        else:
            top = heap[0] # index 0 represent the smallest element

            # if current num is greater than the smallest element (top) then replace top by num in the heap
            # heap will rearrange automatically to keep the new smallest element at index 0
            if num > top:
                heapq.heappushpop(heap, num) # pop followed by a push

    # heap[0] is the smallest element of the heap of length k that contains the larger element of nums ie it is the k-th smallest element
    return heap[0]

Solution 2

The idea of this solution is to use the QuickSelect algorithm as used in QuickSort.

This is the function in QuickSort that finds the correct position of the pivot putting smaller elements on its left and larger elements on its right.

Once QuickSelect is called if the final position of the pivot is k hence return this element, if its position is left to the kth index from end then call the algo on the elements on the right of the pivot (ie larger elements) otherwise call if on the elements on the left of the pivot (ie smaller elements).

# Solution 2
def findKthLargest(nums: List[int], k: int) -> int:
    def quickselect(left, right, nums):
        # Last element will be the pivot and the first element the initial delimiter
        pivot, delimiter = nums[right], left

        for i in range(left, right):
            if nums[i] <= pivot:
                # Swapping values smaller than the pivot to the front
                nums[i], nums[delimiter] = nums[delimiter], nums[i] # 'nums[i]' is now at the index 'delimiter'
                delimiter += 1 # 'nums[i]' is now on the left of the delimiter

        # Finally swapping the last element with the pointer indexed number
        nums[delimiter], nums[right] = nums[right], nums[delimiter]
        return delimiter

    # k_smallest is the index of the kth largest element, left and right the initial limits of the list
    k_smallest, left, right = len(nums) - k, 0, len(nums) - 1
    pivot = quickselect(left, right, nums)

    while pivot != k_smallest:
        # if the current pivot is too small then check the elements at its right
        if pivot < k_smallest:
            left = pivot + 1 # set the left limit as the first element greater than the pivot
            pivot = quickselect(left, right, nums)
        # if the current pivot is too large then check the elements at its right
        else:
            right = pivot - 1  # set the right limit as the first element smaller than the pivot
            pivot = quickselect(left, right, nums)
            
    return(nums[k_smallest])

# Solution 2 bis
def findKthLargest(nums, k):
    if not nums: return

    pivot = random.choice(nums)
    left =  [x for x in nums if x > pivot]
    mid  =  [x for x in nums if x == pivot]
    right = [x for x in nums if x < pivot]

    L, M = len(left), len(mid)

    if k <= L:
        return self.findKthLargest(left, k)
    elif k > L + M:
        return self.findKthLargest(right, k - L - M)
    else:
        return mid[0]