🕌

NeetCode 150 [Heap / Priority Queue]:easy 1/2

2025/02/10に公開

NeetCodeのSolutionを書いていく

Kth Largest Element In a Stream

最初に番号を定めます。数字を足していったときに、最初に指定した番号番目に大きい数字を出力します。

問題の理解

k番目に大きな整数を見つけるクラスを設計します。整数は重複を含みます。
例えば[1,2,3,3]の中で2番目に大きいのは3です。配列はソートされていません。

次の関数を実装しましょう。

  • constructor(int k, int[] nums):kとnumsを使ってオブジェクトを初期化します。
  • int add(int val):valをリストに追加し、k番目に大きな整数を返します。

例1:
Input

["KthLargest",[3,[1,2,3,3]], "add",[3], "add",[5],"add",[6],"add",[7],"add",[8]]

Output:

[null, 3, 3, 3, 5, 6]

Explanation:
KthLargest kthLargest = new KthLargest(3, [1, 2, 3, 3]);
kthLargest.add(3);   // return 3
kthLargest.add(5);   // return 3
kthLargest.add(6);   // return 3
kthLargest.add(7);   // return 5
kthLargest.add(8);   // return 6

処理時間:O(mlogk) mはaddの回数。kは追跡対象が何番目に大きいかを示します。
メモリ:O(k)

解く

求められている処理時間がだいぶ短いな。
そしてデータはソートされていない。
都度ソートしても、クイックソートがO(nlogn)だから大きな問題はないのかな?
ソートに回数が掛けられるからO(k*mlogk)になってしまうからだめか。

わからんので回答見ちゃった・・・

どうやら初めて聞いたがheapという優先度付きキューを使うといいらしい。
ヒープって言うとメモリしか思いつきませんが。

easyにしては難しいなと思ったら。。。解説でもミディアムくらいじゃない?って言われてた。
heapとか実装できないよ!と思って答えみたら、sortとかheap sortとかpythonの定義済みの関数使ってる。
つまりheap sortという関数がpythonにあるということを知っているかという問いになっている。
問題としてイマイチでは?

教科書におけるヒープ構造とは2点で異なる。

  1. ゼロベースインデクスであること
  2. 最小ヒープであること

最小ヒープということは、今回は最大ヒープを使いたいので、そのままでは使えない。
最大ヒープを使うオプションがないかと思ったけどなさそう。
[-3]のインデックスとかで参照できるんだろうか。

どうだ!

class KthLargest:
    _nums = None

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

    def add(self, val: int) -> int:
        heapq.heappush(self._nums, val)
        print(self._nums)
        return self._nums[-3]

だめ!

heapqというのはソートしてくれるものではなく、あくまでheap[k] <= heap[2*k+1] かつ heap[k] <= heap[2*k+2] となる性質を持つだけらしい。
あと、これだとメモリがO[k]では足りなくなっちゃう。

といううことで最後から3番目とかではなく、最小ヒープのサイズを3になるまでpopして、そこからpush popを繰り返し、常に0番目の要素を返す。のが答えでした。

class KthLargest:
    _nums = None
    _k = None

    def __init__(self, k: int, nums: List[int]):
        self._k = k
        self._nums = nums
        heapq.heapify(self._nums)
        while(len(self._nums) > 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]

Discussion