leetcode.com 703. Kth Largest Element in a Stream

 문제 링크: https://leetcode.com/problems/kth-largest-element-in-a-stream/description/

from typing import List

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

    def add(self, val: int) -> int:
        if len(self.nums) == 0:
            self.nums.append(val)
        elif len(self.nums) > 0   \
            and self.nums[len(self.nums) -1] <= val:
            self.nums.append(val)
        else:
            for idx in range(len(self.nums)):
                if val < self.nums[idx]:
                    self.nums.insert(idx, val)
                    break

        return self.nums[len(self.nums) - self.k]

위와 같이 배열의 처음 부터 비교해서, self.nums가 정렬된 상태를 유지한 채로,

val을 오름 차순으로 정렬했을 때, val이 들어갈 위치에, insert() 메서드를 사용해서

집어 넣는 방식으로 구현해도 이 문제의 정답 판정을 받을 수 있습니다.

하지만, 좀... 속도변에서 다른 풀이 보다 느립니다.

그러면 속도를 좀 빠르게 할 방법을 생각해보면, nums의 길이를 줄이는 방법 밖에 없습니다.

문제에서 요구한 것은 뒤에서 부터 k번째 값을 리턴하는 것입니다.

즉 0 ~ (len(nums) -k -1) 인덱스에 있는 값들은 필요하지 않습니다.

따라서, 꼭 가지고 있어야 하는 값은 k 개 입니다.

아래와 같이 heapq를 사용해서, 필요한 가장큰 3개의  값만 남기고

나머지는 모두 제거 하는 방식으로 구현할 수 있습니다.

따라서, add() 가 호출 될 때, 가지고 있는 nums의 길이도 최대 k 로 유지되어서,

실행속도도 빠르게 구현할 수 있습니다.

import heapq
from typing import List

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

        while len(self.nums) > self.k:
            heapq.heappop(self.nums)

    def add(self, val: int) -> int:
        heapq.heappush(self.nums, val)

        while len(self.nums) > self.k:
            heapq.heappop(self.nums)

        return self.nums[0]

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