Notebook

703. Kth Largest Element in a Stream

Design a class to find the k^th largest element in a stream. Note that it is the k^th largest element in the sorted order, not the k^th distinct element.

Implement KthLargest class:

KthLargest(int k, int[] nums) Initializes the object with the integer k and the stream of integers nums.
int add(int val) Returns the element representing the k^th largest element in the stream.

Example 1:

Input
["KthLargest", "add", "add", "add", "add", "add"]
[[3, [4, 5, 8, 2]], [3], [5], [10], [9], [4]]
Output
[null, 4, 5, 5, 8, 8]

Explanation
KthLargest kthLargest = new KthLargest(3, [4, 5, 8, 2]);
kthLargest.add(3);   // return 4
kthLargest.add(5);   // return 5
kthLargest.add(10);  // return 5
kthLargest.add(9);   // return 8
kthLargest.add(4);   // return 8

Constraints:

1 <= k <= 10⁴
0 <= nums.length <= 10⁴
-10⁴ <= nums[i] <= 10⁴
-10⁴ <= val <= 10⁴
At most 10⁴ calls will be made to add.
It is guaranteed that there will be at least k elements in the array when you search for the k^th element.

Source

Code

In [50]:

from bisect import bisect

class KthLargest():
    """naive solution"""
    
    def __init__(self, k, nums):
        self.k = k
        self.array = sorted(nums)


    def add(self, val):
        """Note: maintaining an array of only length k does not speed up the algo
        and using a deque would be unefficient because of O(n) lookup
        """
        idx = bisect(self.array, val)  # ! bisect requires a sorted list
        self.array.insert(idx, val)
        k = -self.k
        return self.array[-k]

In [54]:

import heapq

class KthLargest:
    """Better solution"""

    def __init__(self, k, nums):
        self.k = k
        self.heap = nums
        heapq.heapify(self.heap)  # convert array into heap
        while len(self.heap) > k:
            heapq.heappop(self.heap)

    def add(self, val):
        """Kth largest element --> maintain MinHeap of K elements and return the smallest
           Kth smallest element --> maintain MaxHeap of K elements and return the largest
        """
        if len(self.heap) < self.k:
            heapq.heappush(self.heap, val)
        elif val > self.heap[0]:
            heapq.heappushpop(self.heap, val)  # heapreplace() would work too
        return self.heap[0]
    
# Note: difference between heappushpop(heap, val) and heapreplace(heap, val)
# heapreplace replace smallest value by val ==> if val < min(heap), we have now an even smaller min(heap)
# heappushpop would push val, and pop it out (since it is the min) --> heap would remained unchanged

Check

In [61]:

kthLargest = KthLargest(3, [4, 5, 8, 2])
result = []
result.append(kthLargest.add(3))
result.append(kthLargest.add(5))
result.append(kthLargest.add(10))
result.append(kthLargest.add(9))
result.append(kthLargest.add(4))
print(result)
print(kthLargest.heap)

[4, 4, 4, 4, 4]
[-4, -2, -3]

In [60]:

kthLargest = KthLargest(10,[-10,1,3,1,4,10,3,9,4,5,1])
result = []
result.append(kthLargest.add(3))
result.append(kthLargest.add(2))
result.append(kthLargest.add(3))
result.append(kthLargest.add(1))
result.append(kthLargest.add(2))
result.append(kthLargest.add(4))
result.append(kthLargest.add(5))
result.append(kthLargest.add(5))
result.append(kthLargest.add(6))
result.append(kthLargest.add(7))
result.append(kthLargest.add(7))
result.append(kthLargest.add(8))
result.append(kthLargest.add(2))
result.append(kthLargest.add(3))
result.append(kthLargest.add(1))
result.append(kthLargest.add(1))
result.append(kthLargest.add(1))
result.append(kthLargest.add(10))
result.append(kthLargest.add(11))
result.append(kthLargest.add(5))
result.append(kthLargest.add(6))
result.append(kthLargest.add(2))
result.append(kthLargest.add(4))
result.append(kthLargest.add(7))
result.append(kthLargest.add(8))
result.append(kthLargest.add(5))
result.append(kthLargest.add(6))
print(result)
print(kthLargest.heap)

[5, 4, 4, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2]
[-2, -2, -1, -1, -1, -1, 10, -1, -1, -1]

Follow up:

Return the Kth smallest element instead of the Kth largest

In [57]:

import heapq

class KthLargest:

    def __init__(self, k, nums):
        """only min heap in python --> -nums"""
        self.k = k
        self.heap = [-n for n in nums]
        heapq.heapify(self.heap)  # convert array into heap
        while len(self.heap) > k:
            heapq.heappop(self.heap)

    def add(self, val):
        """Kth smallest element --> maintain MaxHeap of K elements and return the largest"""
        if len(self.heap) < self.k:
            heapq.heappush(self.heap, -val)
        elif val > self.heap[0]:
            heapq.heappushpop(self.heap, -val)
        return -self.heap[0]