Update priority_queue.py

LeeTaek2T · Mar 16, 2024 · 4bd8572 · 4bd8572
1 parent 7a0c5db
commit 4bd8572
Showing 1 changed file with 9 additions and 7 deletions.
diff --git a/priority_queue.py b/priority_queue.py
@@ -8,15 +8,17 @@ def push(self, item, priority):
     self._sift_up(len(self.heap) - 1)
 
   def pop(self):
-    if len(self.heap) > 1:
-      entry = self.heap.pop()
-      self._sift_down(0)
-      return entry[1]
-    elif len(self.heap) == 1:
-      entry = self.heap.pop()
-      return entry[1]
+    if self.heap:
+      if len(self.heap) > 1:
+        top_element = self.heap[0]
+        self.heap[0] = self.heap.pop()
+        self._sift_down(0)
+        return top_element[1]
+      else:
+        return self.heap.pop()[1]
     else:
       return None
+
   def _sift_up(self, index):
     while index > 0:
       parent_index = (index - 1) // 2