50. Kth Largest Element in a Stream(Heap)

문제

k번째로 큰 요소를 찾을 수 있는 클래스를 구현하시오.

KthLargest(int k, int[] nums)로 클래스 생성.

 

def add(int val)은 val을 nums에 추가하고,

val이 추가된 nums의 k번째 요소를 반환 

 

해결 방식(Leetcode 솔루션)

https://leetcode.com/problems/kth-largest-element-in-a-stream/discuss/148866/Python-simple-heapq-solution-beats-100

파이썬의 내장 모듈 heapq 사용.

heapq는 리스트를 최소 힙처럼 다룰 수 있게 해주는 모듈.

import heapq

class KthLargest:
    def __init__(self, k, nums):
        self.k = k
        self.nums = nums
        # nums 리스트를 최소 heap으로 만들기
        heapq.heapify(self.nums)

        # 가장 큰 요소가 -1번 인덱스에 위치하도록,
        # k번째로 큰 요소가 0번 인덱스에 위치하도록 만들기
        while len(self.nums) > k:
            heapq.heappop(self.nums)

    def add(self, val):
        heapq.heappush(self.nums, val)
        # 가장 큰 요소가 -1번 인덱스에 위치하도록,
        # k번째로 큰 요소가 0번 인덱스에 위치하도록 만들기
        if len(self.nums) > self.k:
            heapq.heappop(self.nums)
            
        return self.nums[0]