[leetcode] 703. Kth Largest Element in a Stream 파이썬 코드

https://leetcode.com/problems/kth-largest-element-in-a-stream/description/

 

stream 안에서 k번째 큰 숫자를 찾는 문제인데 

stream이라는 건 숫자를 계속해서 추가할 수 있다는 의미다.

 

이걸 효율적으로 찾을 수 있도록 하는 자료구조가 있는가?

있다. 바로 Min Heap of size K.

 

Min heap에서 log(N) 만에 add or pop 할 수 있다.

최솟값은 O(1) 만에 찾을 수 있다.

 

Max Heap 대신에 Min heap을 사용하는 이유는?

 

여기서는 값을 remove 하지 않고 추가만 한다. 

[4, 5, 8, 2]라는 숫자 배열이 있고 k가 3일 때 

while the size of the loop is greater than k, pop minimum value.

여기서는 2가 제일 작으니까 2를 pop하고 나면, [4, 5, 8]이 남고 여기서 3번째로 큰 수는 4이다.

min heap에서 최솟값을 보는 것은 O(1)이므로 매우 효율적이다.

 

class KthLargest(object):

    def __init__(self, k, nums):
        self.minHeap, self.k = nums, k
        heapq.heapify(self.minHeap)
        while len(self.minHeap) > k: 
            heapq.heappop(self.minHeap)

    def add(self, val):
        heapq.heappush(self.minHeap, val)
        if len(self.minHeap) > self.k:
            heapq.heappop(self.minHeap)
        return self.minHeap[0]