LeetCode/NeetCode

Heap-PrioiryQueue: 703. Kth Largest Element in a Stream ★

hyunkookim 2025. 1. 20. 13:49

703. Kth Largest Element in a Stream

 

https://youtu.be/hOjcdrqMoQ8?si=j-8DMVDjR9t8bNx_

 

class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        self.k = k
        self.minHeap = nums
        heapq.heapify(self.minHeap)

        # 최소 힙의 크기를 k 개로 유지하기 위한 로직
        while len(self.minHeap) > k:
            # nums 리스트의 길이가 k보다 큰 경우, 
            # 가장 작은 값을 하나씩 제거하여 힙의 크기를 k로 줄임
            heapq.heappop(self.minHeap)

        """
        KthLargest 클래스는 항상 k번째로 큰 값을 관리
        최소 힙은 항상 루트 노드에 가장 작은 값을 유지
        힙의 크기를 k로 제한하면, 
        힙의 루트 노드(self.minHeap[0])는 k번째로 큰 값이 됨

        k = 3
        nums = [4, 5, 8, 2]

        1. heapq.heapify(nums) -> 최소 힙 생성: [2, 4, 5, 8]
        2. 크기 조정: 현재 힙 크기: 4 > 3 
            -> 가장 작은 값 2 제거하여 힙: [4,5,8]
        3. 결과: self.minHeap[0] = 4 최소값: 3 번째로 큰 값
        """

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


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

 

최종 코드

class KthLargest:

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

        # k번째 큰값 찾는 것이니, k+1번째 큰 값부터는 필요없음
        # 즉, 작은 값들은 필요없음
        # 따라서, 최소 heap으로 만들고
        heapq.heapify(self.minHeap)
        # 개수는 k개로 유지, pop 하면 가장 작은값 삭제
        while(len(self.minHeap) > k): # k개 보다 크면 작은거 삭제
            heapq.heappop(self.minHeap)
        # 그럼, 끝 값이 k 번째 큰값임

    def add(self, val: int) -> int:
        # val 추가
        heapq.heappush(self.minHeap, val)

        # 개수가 k 보다 크면, 하나 삭제
        if len(self.minHeap) > self.k:
            heapq.heappop(self.minHeap)

        # 총 개수가 k 개 이므로, min heap 이어서, 
        # 가장 작은 값 self.minHeap[0] 는 k번째 큰값임
        return self.minHeap[0]
        


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