stream이 주어질 때 순서대로 k번째로 큰 수를 출력하는 문제

Solution

min heap으로 간단히 구현할 수 있다.
처음에는 heap 2개를 썼는데 k개만 관리하면 돼서 k 길이의 min heap 하나면 충분하다.

k 길이의 min heap을 관리하면서 min heap에서 가장 작은 값보다 큰 값이 들어오면 min heap에서 빼고 갱신해주면 된다.

Source Code

class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        nums.sort()

        self.k = k
        self.largeHeap = nums[-k:]
        heapq.heapify(self.largeHeap)

    def add(self, val: int) -> int:
        if len(self.largeHeap) < self.k:
            heapq.heappush(self.largeHeap, val)
        elif val > self.largeHeap[0]:
            heapq.heappop(self.largeHeap)
            heapq.heappush(self.largeHeap, val)

        return self.largeHeap[0]


# Your KthLargest object will be instantiated and called as such:
# obj = KthLargest(k, nums)
# param_1 = obj.add(val)

703. Kth Largest Element in a Stream

Problem link

Problem Summary

Solution

Source Code