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

NeetCodeのSolutionを書いていく

Last Stone Weight

stonesという整数の配列が与えられます。stone[i]は1番目の石の重さを表します。
石を使ったシミュレーションをしたいです。

• 各ステップでもっと重い2つの石x,yを選びます。それぞれぶつけ合わせてバラバラにします。
• もしxとyが同じであれば両方の石が砕けます。
• もしxがyより小さければ重さxの石が砕けます。そして重さyの石は削られて重さがy-xになります。

残り一つになるまでシミュレーションを続けます。
最後に残った石の重さかあるいはなければ0を返します。

制約

石の数は1個以上20個以下
石の重さは1以上100以下

期待される時間及びメモリ

時間：O(nlong)
メモリ：O(n)
nは石の数

解く

Kth Largest Element In a Streamを参照にすると、heapifyを使うわけだよね。
今回は石の重さは1以上だから、最小値ヒープで最大値を扱うには、一時的に負数にしてやればいいね。
それで大きい方からぶつけていけばいい。

どうだ！

class Solution:
    def lastStoneWeight(self, stones: List[int]) -> int:
        heap = [-w for w in stones]
        heapq.heapify(heap)
        d = 0
        while(len(heap)>=2):
            print(heap)
            d = heap[0] - heap[1]
            heapq.heappop(heap)
            heapq.heappop(heap)
            if d != 0:
                heapq.heappush(heap, d)
        if len(heap) == 0:
            return 0
        return -heap[0]

だめ！

そういえば最小値ヒープはソートしてくれるものではなかった・・・
なので [7,5,8] というパターンで [8,7,5] となってくれるとは限らず [8,5,7] となる場合もあり、これだと配列の0番目は8、1番目は5なので計算できない。
ので、一個ずつ取り出してやればいい！

class Solution:
    def lastStoneWeight(self, stones: List[int]) -> int:
        heap = [-w for w in stones]
        heapq.heapify(heap)
        d = -1
        while(len(heap)>=2):
            print(heap)
            a = heap[0]
            heapq.heappop(heap)
            b = heap[0]
            heapq.heappop(heap)
            d = a - b
            if d != 0:
                heapq.heappush(heap, d)
        if len(heap) == 0:
            return 0
        return -heap[0]

いけた！