refactor PriorityQueue and tests (keon#348)

goswami-rahul · Hai Hoang Dang · commit 1ac287f4a84c · 2018-06-10T10:55:31.000-07:00
diff --git a/algorithms/queues/priority_queue.py b/algorithms/queues/priority_queue.py
@@ -3,37 +3,54 @@
 Insertion - O(n)
 Extract min/max Node - O(1)
 """
-import collections
+import itertools
 
 
 class PriorityQueueNode:
-	def __init__(self, data, priority):
-		self.data = data
-		self.priority = priority
+    def __init__(self, data, priority):
+        self.data = data
+        self.priority = priority
+
+    def __repr__(self):
+        return "{}: {}".format(self.data, self.priority)
 
-	def __repr__(self):
-		return str(self.data) + ": " + str(self.priority)
 
 class PriorityQueue:
-	def __init__(self):
-		self.priority_queue_list = collections.deque()
-
-	def __repr__(self):
-	    return "PriorityQueue({!r})".format(list(self.priority_queue_list))
-
-	def size(self):
-		return len(self.priority_queue_list)
-
-	def push(self, item, priority=None):
-		priority = item if priority is None else priority
-		node = PriorityQueueNode(item, priority)
-		for index, current in enumerate(self.priority_queue_list):
-			if current.priority > node.priority:
-				self.priority_queue_list.insert(index, node)
-				return
-		# when traversed complete queue
-		self.priority_queue_list.append(node)
-
-	def pop(self):
-		# remove and return the first node from the queue
-		return self.priority_queue_list.popleft()
+    def __init__(self, items=None, priorities=None):
+        """Create a priority queue with items (list or iterable).
+        If items is not passed, create empty priority queue."""
+        self.priority_queue_list = []
+        if items is None:
+            return
+        if priorities is None:
+            priorities = itertools.repeat(None)
+        for item, priority in zip(items, priorities):
+            self.push(item, priority=priority)
+
+    def __repr__(self):
+        return "PriorityQueue({!r})".format(self.priority_queue_list)
+
+    def size(self):
+        """Return size of the priority queue.
+        """
+        return len(self.priority_queue_list)
+
+    def push(self, item, priority=None):
+        """Push the item in the priority queue.
+        if priority is not given, priority is set to the value of item.
+        """
+        priority = item if priority is None else priority
+        node = PriorityQueueNode(item, priority)
+        for index, current in enumerate(self.priority_queue_list):
+            if current.priority < node.priority:
+                self.priority_queue_list.insert(index, node)
+                return
+        # when traversed complete queue
+        self.priority_queue_list.append(node)
+
+    def pop(self):
+        """Remove and return the item with the lowest priority.
+        """
+        # remove and return the first node from the queue
+        return self.priority_queue_list.pop().data
+
diff --git a/tests/test_queues.py b/tests/test_queues.py
@@ -1,11 +1,12 @@
+import unittest
+
 from algorithms.queues import (
     ArrayQueue, LinkedListQueue,
     max_sliding_window,
     reconstruct_queue,
-    PriorityQueue, PriorityQueueNode
+    PriorityQueue
 )
 
-import unittest
 
 class TestQueue(unittest.TestCase):
     """
@@ -70,10 +71,9 @@ def test_LinkedListQueue(self):
 
         self.assertTrue(queue.is_empty())
 
-class TestSuite(unittest.TestCase):
 
+class TestSuite(unittest.TestCase):
     def test_max_sliding_window(self):
-
         array = [1, 3, -1, -3, 5, 3, 6, 7]
         self.assertEqual(max_sliding_window(array, k=5), [5, 5, 6, 7])
         self.assertEqual(max_sliding_window(array, k=3), [3, 3, 5, 5, 6, 7])
@@ -85,24 +85,23 @@ def test_max_sliding_window(self):
         self.assertEqual(max_sliding_window(array, k=2), [8, 10, 10, 9, 9, 15, 15, 90, 90])
 
     def test_reconstruct_queue(self):
-        self.assertEqual([[5,0], [7,0], [5,2], [6,1], [4,4], [7,1]],
-                         reconstruct_queue([[7,0], [4,4], [7,1], [5,0], [6,1], [5,2]]))
+        self.assertEqual([[5, 0], [7, 0], [5, 2], [6, 1], [4, 4], [7, 1]],
+                         reconstruct_queue([[7, 0], [4, 4], [7, 1], [5, 0], [6, 1], [5, 2]]))
 
-"""
-    TODO: Refactor PriorityQueue because insert method does not work for python3.4
-class TestPriorityQueue(unittest.TestCase):
 
-        Test suite for the PriorityQueue data structures.
+class TestPriorityQueue(unittest.TestCase):
+    """Test suite for the PriorityQueue data structures.
+    """
 
     def test_PriorityQueue(self):
-        queue = PriorityQueue()
-        queue.push(3)
-        queue.push(4)
-        queue.push(1)
-        queue.push(6)
-        self.assertEqual(4,queue.size())
-        self.assertEqual(str(1) + ": " + str(1),str(queue.pop()))
-"""
-if __name__ == "__main__":
+        queue = PriorityQueue([3, 4, 1, 6])
+        self.assertEqual(4, queue.size())
+        self.assertEqual(1, queue.pop())
+        self.assertEqual(3, queue.size())
+        queue.push(2)
+        self.assertEqual(4, queue.size())
+        self.assertEqual(2, queue.pop())
+
 
+if __name__ == "__main__":
     unittest.main()
