Update utils4e.py

utils4e.py

@@ -21,62 +21,76 @@
 
 
 class PriorityQueue:
-    """A Queue in which the minimum (or maximum) element (as determined by f and order) is returned first.
-    If order is 'min', the item with minimum f(x) is
-    returned first; if order is 'max', then it is the item with maximum f(x).
+    """A Queue in which the minimum (or maximum) element 
+    (as determined by f and order) is returned first.
+    If order is 'min', then item with minimum f(x) is returned first;
+    if the order is 'max', then it is the item with maximum f(x).
     Also supports dict-like lookup."""
 
     def __init__(self, order='min', f=lambda x: x):
         self.heap = []
-
+        self.items = {}
+        self.length = 0
         if order == 'min':
             self.f = f
-        elif order == 'max':  # now item with max f(x)
-            self.f = lambda x: -f(x)  # will be popped first
+        elif order == 'max':
+            self.f = lambda x: -f(x)    # now item with max f(x) will be popped first.
         else:
             raise ValueError("Order must be either 'min' or 'max'.")
 
     def append(self, item):
         """Insert item at its correct position."""
         heapq.heappush(self.heap, (self.f(item), item))
+        if item not in self.items:
+            self.items[item] = [1, self.f(item)]
+        else:
+            self.items[item][0] += 1
+        self.length += 1
 
     def extend(self, items):
         """Insert each item in items at its correct position."""
         for item in items:
             self.append(item)
 
     def pop(self):
-        """Pop and return the item (with min or max f(x) value)
-        depending on the order."""
-        if self.heap:
-            return heapq.heappop(self.heap)[1]
+        """Pop and return the item (with min or max f(x) value) depending
+        on the order."""
+        while self.heap:
+            item = heapq.heappop(self.heap)[1]
+            if item not in self.items:
+                continue
+            if self.items[item][0] > 1:
+                self.items[item][0] -= 1
+            elif self.items[item][0] == 1:
+                del self.items[item]
+            return item
+            self.length -= 1
+
         else:
-            raise Exception('Trying to pop from empty PriorityQueue.')
+            raise Exception("Trying to pop from empty PriorityQueue.")
 
     def __len__(self):
         """Return current capacity of PriorityQueue."""
-        return len(self.heap)
+        return self.length
 
     def __contains__(self, key):
-        """Return True if the key is in PriorityQueue."""
-        return any([item == key for _, item in self.heap])
+        """Return True if key is in PriorityQueue."""
+        return key in self.items
 
     def __getitem__(self, key):
-        """Returns the first value associated with key in PriorityQueue.
-        Raises KeyError if key is not present."""
-        for value, item in self.heap:
-            if item == key:
-                return value
-        raise KeyError(str(key) + " is not in the priority queue")
+        """Returns the value associated with key in PriorityQueue.
+        Raise KeyError if key is not present."""
+        if key in self.items:
+            return self.items[key][1]
+        raise KeyError(str(key) + ' is not in the priority queue')
 
     def __delitem__(self, key):
         """Delete the first occurrence of key."""
         try:
-            del self.heap[[item == key for _, item in self.heap].index(True)]
-        except ValueError:
+            del self.items[key]
+            self.length -= 1
+        except KeyError:
             raise KeyError(str(key) + " is not in the priority queue")
-        heapq.heapify(self.heap)
-
 
 # ______________________________________________________________________________
 # Functions on Sequences and Iterables