heap updated

Author: shaonty

Data Structure/Heap/heap.py

@@ -3,75 +3,52 @@
 
 class MaxHeap(object):
 
-	HEAP_SIZE = 10
-
-	def __init__(self):
-		self.heap = [0]*MaxHeap.HEAP_SIZE
-		self.currentPosition = -1
+	def __init__(self, maxSize=None):
+		self.heap = [float("inf")] # heap starting from index 1
+		self.HEAP_SIZE = maxSize
 
 	def _swap(self,i,j):
 		self.heap[i], self.heap[j] = self.heap[j], self.heap[i]
 
 	def _heapIsFull(self):
-		if self.currentPosition == MaxHeap.HEAP_SIZE:
-			return True
-		return False
+		return (self.HEAP_SIZE != None and len(self.heap) >= self.HEAP_SIZE)
 
 	def insert(self, item):
 		if self._heapIsFull():
 			"Heap is full..."
-			return
-		self.currentPosition+=1
-		self.heap[self.currentPosition] = item
-		# adjust parent node item
-		self._bubbleUp(self.currentPosition)
+		else:
+			self.heap.append(item)
+			# adjust parent node item
+			self._bubbleUp(len(self.heap)-1)
 
 	def _bubbleUp(self, currentPosition):
-		index = currentPosition
-		parrentIndex = index//2
+		if currentPosition >= 2: # no need to do bubbleUp for 1 element
+			index = currentPosition
+			parrentIndex = index//2
 
-		if parrentIndex>=0 and self.heap[parrentIndex] < self.heap[index]:
-			self._swap(parrentIndex, index)
-			# index = parrentIndex
-			self._bubbleUp(parrentIndex)
+			if parrentIndex >= 0 and self.heap[parrentIndex] < self.heap[index]:
+				self._swap(parrentIndex, index)
+				self._bubbleUp(parrentIndex)
 
-
-	def getMax(self):
-		if self.heap[0]:
-			return self.heap[0]
-		else:
-			return "Empty Heap"
+	def peek(self):
+		return self.heap[1] if len(self.heap) > 1 else False
 
 	def pop(self):
-		if self.currentPosition > 1:
-			# swap first element with last element
-			self._swap(0,self.currentPosition)
-			del self.heap[self.currentPosition]
-			# adjust parent node item
-			self._bubbleDown(0)
-		elif self.currentPosition == 1:
-			del self.heap[0]
-		else:
-			print "No deletion operation for empty heap"
-			return "False"
+		element = self.peek()
+		if element:
+			self._swap(1, len(self.heap) - 1)
+			self.heap.pop()
+			self._bubbleDown(1)
+		return element
 
 	def _bubbleDown(self, index):
-		left = 2*index+1
-		right = 2*index+2
+		leftChildIndex = 2 * index
+		rightChildIndex = 2 * index + 1
 		largest = index
-		if len(self.heap) > left and self.heap[largest] < self.heap[left]:
-			largest = left
-		if len(self.heap) > right and self.heap[largest] < self.heap[right]:
-			largest = right
+		if len(self.heap) > leftChildIndex and self.heap[largest] < self.heap[leftChildIndex]:
+			largest = leftChildIndex
+		if len(self.heap) > rightChildIndex and self.heap[largest] < self.heap[rightChildIndex]:
+			largest = rightChildIndex
 		if largest!=index:
 			self._swap(index, largest)
 			self._bubbleDown(largest)
-
-	def printHeap(self):
-		all_hp_element = []
-		for i in self.heap:
-			if i!=0:
-				all_hp_element.append(i)
-		return all_hp_element
-
-			

Data Structure/Heap/test_heap.py

@@ -12,13 +12,18 @@ def test(self):
         ob.insert(3)
         ob.insert(8)
         ob.insert(20)
-
+        self.assertEqual(ob.peek(), 20, msg="Max Element is not matched")
+        ob.pop()
+        ob.pop()
+        self.assertEqual(ob.peek(), 8, msg="Max Element is not matched")
+        ob.pop()
+        self.assertEqual(ob.peek(), 6, msg="Max Element is not matched")
         ob.pop()
         ob.pop()
-        
-        get_max = ob.getMax()
-        self.assertEqual(ob.printHeap(), [8, 6, 3, 5], msg="Heap Element Should be [8,6,3,5]")
-        self.assertEqual(get_max, 8, msg="Max Element is not matched")
+        ob.pop()
+        self.assertEqual(ob.peek(), False, msg="Max Element is not matched")
 
 if __name__ == '__main__':
     unittest.main()
+
+