from treenode import TreeNode
class BinarySearchTree:
def __init__(self):
self.root = None
self.size = 0
def length(self):
return self.size
def __len__(self):
return self.size
def __iter__(self):
return self.root.__iter__()
def get(self,key):
if self.root is not None:
res = self._get(key,self.root)
if res is not None:
return res.payload
else:
return None
else:
return None
def _get(self,key,currentNode):
if not currentNode:
return None
elif currentNode.key == key:
return currentNode
elif key < currentNode.key:
return self._get(key,currentNode.leftChild)
else:
return self._get(key,currentNode.rightChild)
def _put(self,key,val,currentNode):
"""helper function for put(), finds the correct
placing for a new node"""
if key < currentNode.key:
if currentNode.hasLeftChild():
self._put(key,val,currentNode.leftChild)
else:
currentNode.leftChild = TreeNode(key,val,parent=currentNode)
else:
if currentNode.hasRightChild():
self._put(key,val,currentNode.rightChild)
else:
currentNode.rightChild = TreeNode(key,val,parent=currentNode)
def put(self,key,val):
node = self._get(key,self.root)
if node is None:
if self.root:
self._put(key,val,self.root)
else:
self.root = TreeNode(key,val)
self.size += 1
else:
#key is already in tree
#replace the value
node.payload = val
def __getitem__(self, key):
return self.get(key)
def remove(self, currentNode):
"""remove a node from the tree"""
if currentNode.isLeaf():
if currentNode == currentNode.parent.leftChild:
currentNode.parent.leftChild = None
else:
currentNode.parent.rightChild = None
elif currentNode.hasBothChildren():
succ = currentNode.findSuccessor()
succ.spliceOut()
currentNode.key = succ.key
currentNode.payload = succ.payload
else:
if currentNode.hasLeftChild():
if currentNode.isLeftChild():
currentNode.leftChild.parent = currentNode.parent
currentNode.parent.leftChild = currentNode.leftChild
elif currentNode.isRightChild():
currentNode.leftChild.parent = currentNode.parent
currentNode.parent.rightChild = currentNode.leftChild
else:
#currentNode is root
currentNode.replaceNodeData(currentNode.leftChild.key,
currentNode.leftChild.payload,
currentNode.leftChild.leftChild,
currentNode.leftChild.rightChild)
else: #has right child
if currentNode.isLeftChild():
currentNode.rightChild.parent = currentNode.parent
currentNode.parent.leftChild = currentNode.rightChild
elif currentNode.isRightChild():
currentNode.rightChild.parent = currentNode.parent
currentNode.parent.rightChild = currentNode.rightChild
else:
currentNode.replaceNodeData(currentNode.rightChild.key,
currentNode.rightChild.payload,
currentNode.rightChild.leftChild,
currentNode.rightChild.rightChild)
def delete(self, key):
if self.size > 1:
nodeToRemove = self._get(key, self.root)
if nodeToRemove is not None:
self.remove(nodeToRemove)
self.size -= 1
else:
raise KeyError("Key is not in tree")
elif self.size == 1 and self.root.key == key:
self.root = None
self.size -= 1
else:
raise KeyError("Cannot delete key from an empty tree")
def __delitem__(self, key):
self.delete(key)
def getAsList(self):
"""Non-recursive inorder traversal"""
result = []
currentNode = None
if self.size == 1:
result.append(str(self.root.key) + ":" + str(self.root.payload))
elif self.size > 1:
currentNode = self.root.findMin()
while currentNode is not None:
result.append(str(currentNode.key) + ":" + str(currentNode.payload))
currentNode = currentNode.findSuccessor()
return result