class BinarySearchTree:
# Part 0: setting up our BST
# my_bst = BinarySearchTree()
def __init__(self):
self.root = None
self.size = 0
# my_bst.length()
def length(self):
return self.size
# len(my_bst)
def __len__(self):
return self.size
# Part 1: Inserting a key/value pair into a BST
# function for external use
def put(self, key, value):
if not self.root:
# creates root if tree is empty
self.root = BSTNode(key, value)
else:
# otherwise call helper function to do the work
self._put(key, value, self.root)
self.size = self.size + 1
# internal helper function for node insertion
# "_" is a convention for internal-use only functions
# recursive
def _put(self, key, value, current_node):
# case 1: go to left subtree
if key < current_node.key:
# go to the subtree recursively
if current_node.left_child:
self._put(key, value, current_node.left_child)
# or add a leaf if no left_child (base case)
else:
current_node.left_child = BSTNode(key, value, parent = current_node)
# case 2: found the key, do an update (base case)
elif key == current_node.key:
current_node.value = value
# case 3: go right (key is bigger than current node's key)
else:
if current_node.right_child:
self._put(key, value, current_node.right_child)
else:
current_node.right_child = BSTNode(key, value, parent = current_node)
# lets us do my_bst[key] = value
def __setitem__(self, key, value):
self.put(key, value)
# Part 2: Getting a value for a given key in a BST
def get(self, key):
# if tree is not empty
if self.root:
# use helper function to get the node with that key
result_node = self._get(key, self.root)
# if result_node is not None, we found the key
if result_node:
return result_node.value
# if result_node is None, then we didn't find the key
else:
return None
# if tree is empty, return None
else:
return None
# get helper function
# recursive function, either returns node with matching key
# or returns None if key is not in the bst
def _get(self, key, current_node):
# two base cases, not found and found
if current_node is None:
return None
elif current_node.key == key:
return current_node
# make one of two recursive calls
elif current_node.key > key:
# search left subtree
return self._get(key, current_node.left_child)
else:
# search right subtree
return self._get(key, current_node.right_child)
# lets us do print(my_bst[key]) and get the value, like a dictionary
def __getitem__(self, key):
return self.get(key)
# Part 3: Check if a key is in a BST
def __contains__(self, key):
if self._get(key, self.root) is not None:
return True
return False
# Part 4: Delete a node from a BST
def delete(self, key):
if self.size > 1:
node_to_remove = self._get(key, self.root)
if node_to_remove:
self.remove(node_to_remove)
self.size = self.size - 1
else:
raise KeyError
elif self.size == 1 and self.root.key == key:
self.root = None
self.size = self.size - 1
else:
raise KeyError('Error: Key is not in tree')
# lets us use "del" operator
def __delitem__(self, key):
self.delete(key)
# does the actual work of removing a node
# the tree will need to be restructered and there are multiple cases
def remove(self, current_node):
# Case 1: removing a node with no children
if current_node.is_leaf():
if current_node.is_left_child():
current_node.parent.left_child = None
else:
current_node.parent.right_child = None
# Case 2: removing a node with only one child
elif not current_node.has_both_children():
# if the child it has is the left_child
if current_node.left_child:
# A
if current_node.is_left_child():
current_node.left_child.parent = current_node.parent
current_node.parent.left_child = current_node.left_child
# C
elif current_node.is_right_child():
current_node.right_child.parent = current_node.parent
current_node.parent.right_child = current_node.right_child
# root with left child
else:
current_node.replace_node_data(current_node.left_child.key, \
current_node.left_child.value, \
current_node.left_child.left_child, \
current_node.left_child.right_child)
# do the same thing except for right_child
else:
# B
if current_node.is_left_child():
current_node.right_child.parent = current_node.parent
current_node.parent.left_child = current_node.right_child
# D
elif current_node.is_right_child():
current_node.right_child.parent = current_node.parent
current_node.parent.right_child = current_node.right_child
# root with right child
else:
current_node.replace_node_data(current_node.right_child.key, \
current_node.right_child.value, \
current_node.right_child.left_child, \
current_node.right_child.right_child)
# Case 3: removing a node with both children
else:
successor = current_node.find_successor()
successor.splice_out()
current_node.key = successor.key
current_node.value = successor.value
# Part 5: Iterate over a BST
def __iter__(self):
# calls __iter__ in BSTNode
return self.root.__iter__()
