[NeetCode/Python] Kth Largest Integer in a Stream

문제: https://neetcode.io/problems/kth-largest-integer-in-a-stream

NeetCode

 

문제에서 제공하는대로 리스트에 숫자를 더 추가하고, 매 리스트에서 k번째로 가장 큰 숫자를 리턴하는 것이다.

 

[내 코드]

class KthLargest:

    def __init__(self, k: int, nums: List[int]):
        self.k = -k
        self.nums = sorted(nums)
        return

    def add(self, val: int) -> int:
        self.nums.append(val)
        self.nums = sorted(self.nums)
        return self.nums[self.k]

 

정말 단순한 풀이.. add를 시킬 때마다 정렬을 하고, 그 리스트에서 k번째 숫자를 리턴한다.

시간 복잡도가 상당한 풀이라 더 좋은 풀이가 없나 고민하면 좋았을 텐데 우선 푸는것에 집중했다 😅 (스터디 마감 시간 다돼서 촉박하게 풀었다..)

 

[솔루션 코드]

class KthLargest:
    
    def __init__(self, k: int, nums: List[int]):
        self.minHeap, self.k = nums, k
        heapq.heapify(self.minHeap)
        while len(self.minHeap) > k:
            heapq.heappop(self.minHeap)

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

 

heap을 쓰면 된다니!!! 자료구조에 다시 익숙해져야 한다!! 최댓값, 최솟값 찾는건 완전 이진트리, 더 나아가 힙을 바로 떠올려야 한다!!

heapq.heapify : 리스트를 힙으로 변환한다. O(n)

어짜피 k번째로 큰 숫자들만 관리하면 되므로, k개의 숫자만 관리하면 된다.

minHeap으로 관리하면서 k개를 맞추기 위해 작은 숫자들은 pop해주면, root값이 바로 k번째로 큰 숫자가 된다!

(maxHeap을 만들고 싶다면 음수로 숫자를 관리하면 될 듯)

 

천재... 진짜 알고리즘은 풀면서 자극 받는 이런 맛이 있다.. 실무하면서 굳어진 뇌를 말랑말랑하게 풀어가는 중이다!

 

시간 복잡도

n은 nums의 길이

-> 내 풀이: O(n)

-> 솔루션: O(logn), init t시에는 O(n)