Heaps: kth largest element

A heap (priority queue) always hands you the smallest item first, in O(log n) per push/pop. Python's heapq turns a plain list into a min-heap:

import heapq
h = []
heapq.heappush(h, 5)
heapq.heappush(h, 2)
h[0]                 # 2     — the smallest is always at the front
heapq.heappop(h)     # 2     — remove and return the smallest

This is LeetCode #215. To find the kth largest with a min-heap, keep a heap of size k: the smallest of the k biggest values you have seen sits at the top, and that is exactly the kth largest once you have scanned everything.

Your task

Write kth_largest(nums, k) returning the kth largest element (1-indexed, duplicates count). So k = 1 is the maximum.

kth_largest([3, 2, 1, 5, 6, 4], 2)            # 5
kth_largest([3, 2, 3, 1, 2, 4, 5, 5, 6], 4)   # 4
kth_largest([1], 1)                            # 1

Hint: push each number onto a min-heap; whenever it grows past size k, pop the smallest. The element left at heap[0] is your answer. (heapq.nlargest also works if you prefer.)