Implementation of min/max heap

heap.py

import math

"""Implementation of min heap
Works with Python 2 or 3
Run from terminal:
 `python heap.py`
"""

class Heap(object):

	"""Create a new Heap object
	argument data (Number): first element/top of the heap
	"""
	def __init__(self,data,min_or_max="min"):
		self.arr = [data]
		self.min_or_max = min_or_max

	"""Establish correct min/max function for this Heap
	argument a, b (Number): elements to compare
	"""
	def minOrMax(self,a,b):
		if self.min_or_max == "min":
			return min(a,b)
		else:
			return max(a,b)

	"""Establish correct parent-child relationship for this Heap
	argument p, c (Number): parent and child to compare
	"""
	def parentChild(self,p,c):
		if self.min_or_max == "min":
			return (p < c)
		else:
			return (p > c)

	"""Recursively look up the heap to move smaller values up
	argument pos (Number): position from which to begin sifting
	"""
	def siftUp(self,pos):
		if pos == 0:
			return
		child = self.arr[pos]
		parentpos = int(math.floor((pos-1)/2))
		parent = self.arr[parentpos]
		if not self.parentChild(parent,child):
			self.arr[parentpos] = child
			self.arr[pos] = parent
			self.siftUp(parentpos)

	"""Recursively look down the heap to move smaller values up
	argument pos (Number): position from which to begin sifting
	"""
	def siftDown(self,pos):
		parent = self.arr[pos]
		## no children
		if len(self.arr) < (2*pos + 2):
			return
		## left child only
		elif len(self.arr) < (2*pos + 3):
			lchild = self.arr[2*pos + 1]
			if not self.parentChild(parent,lchild):
				self.arr[pos] = lchild
				self.arr[2*pos + 1] = parent
				self.siftDown(2*pos + 1)
		## both children; compare to prioritize
		else:
			lchild = self.arr[2*pos + 1]
			rchild = self.arr[2*pos + 2]
			if not self.parentChild(parent,self.minOrMax(lchild,rchild)):
				if self.parentChild(lchild,rchild):
					self.arr[pos] = lchild
					self.arr[2*pos + 1] = parent
					self.siftDown(2*pos + 1)
				else:
					self.arr[pos] = rchild
					self.arr[2*pos + 2] = parent
					self.siftDown(2*pos + 2)

	"""Insert a new element/leaf into the Heap
	argument data (Number or Array): element to add to Heap
	"""
	def insert(self,data):
		if isinstance(data,list):
			for i in data:
				self.arr.append(i)
				self.siftUp(len(self.arr)-1)
		else:
			self.arr.append(data)
			self.siftUp(len(self.arr)-1)

	"""Delete the first element/top from Heap and sift down list
	no arguments
	"""
	def delete(self):
		self.arr[0] = self.arr[len(self.arr)-1]
		self.arr.pop()
		self.siftDown(0)

"""Tests
"""
print("TEST 1")
myMinHeap = Heap(3,"min")
myMinHeap.insert([15,12,21,25])
print('initial')
print(myMinHeap.arr)

myMinHeap.insert(6)
print('insert 6')
print(myMinHeap.arr)

myMinHeap.delete()
print('delete')
print(myMinHeap.arr)

print("TEST 2")
myMaxHeap = Heap(40,"max")
myMaxHeap.insert([15,22,13,8,23])
print('initial')
print(myMaxHeap.arr)

myMaxHeap.insert(45)
print('insert 45')
print(myMaxHeap.arr)

myMaxHeap.delete()
print('delete')
print(myMaxHeap.arr)