O(nlogn) time and O(n) space using max heap.

149ps · Apr 7, 2022 · b8c7c0e · b8c7c0e
1 parent dc8a236
commit b8c7c0e
Showing 1 changed file with 5 additions and 9 deletions.
diff --git a/1046. Last Stone Weight/1046. Last Stone Weight.py b/1046. Last Stone Weight/1046. Last Stone Weight.py
@@ -31,12 +31,8 @@ def lastStoneWeight(self, stones: List[int]) -> int:
         max_heap = [-val for val in stones]
         heapq.heapify(max_heap)
         while len(max_heap) > 1:
-            y = (-1) * heapq.heappop(max_heap)
-            x = (-1) * heapq.heappop(max_heap)
-            if x == y:
-                if not max_heap: # consider a case [1,1]
-                    return 0
-                continue
-            else:
-                heapq.heappush(max_heap,x-y)
-        return max_heap[0] if max_heap[0] > 0 else (-1)*max_heap[0]
+            stone1 = (-1) * heapq.heappop(max_heap)
+            stone2 = (-1) * heapq.heappop(max_heap)
+            if stone1 != stone2:
+                heapq.heappush(max_heap,-(stone1-stone2))
+        return -max_heap[0] if max_heap else 0