130 lines
3.3 KiB
Python
130 lines
3.3 KiB
Python
# Also do the 2026.05.06 homework
|
|
#
|
|
# Implement insertion, removal and search in a binary search tree.
|
|
# Properties of binary search trees:
|
|
# - Each node has at most two children
|
|
# - The left subtree of a node contains only nodes with values less than the node's value
|
|
# - The right subtree of a node contains only nodes with values greater than the node's value
|
|
#
|
|
# Example of what such a tree would look like:
|
|
# 10
|
|
# / \
|
|
# 5 15
|
|
# / \ \
|
|
# 3 7 20
|
|
#
|
|
# Example of an *invalid* tree:
|
|
# 10
|
|
# / \
|
|
# 5 15
|
|
# / \ \
|
|
# 3 7 20
|
|
# \
|
|
# 12
|
|
#
|
|
# Don't worry about balancing the tree for now, we will cover that later
|
|
|
|
|
|
class Node:
|
|
def __init__(self, value, left=None, right=None):
|
|
self.value = value
|
|
self.left: Node | None = left
|
|
self.right: Node | None = right
|
|
|
|
|
|
class BinarySearchTree:
|
|
def __init__(self, root=None):
|
|
self.root: Node | None = root
|
|
|
|
def find_prev_node(self, value):
|
|
prev_node: Node | None = None
|
|
node = self.root
|
|
while node is not None:
|
|
prev_node = node
|
|
if node.value >= value:
|
|
node = node.left
|
|
else:
|
|
node = node.right
|
|
|
|
return prev_node
|
|
|
|
def find_existing_node_and_parent(self, value):
|
|
prev_node: Node | None = None
|
|
node = self.root
|
|
|
|
while node is not None:
|
|
if value == node.value:
|
|
return node, prev_node
|
|
|
|
prev_node = node
|
|
if value < node.value:
|
|
node = node.left
|
|
else:
|
|
node = node.right
|
|
|
|
return None, None
|
|
|
|
def insert(self, value):
|
|
prev_node = self.find_prev_node(value)
|
|
if prev_node is None:
|
|
self.root = Node(value)
|
|
return
|
|
|
|
if prev_node.value >= value:
|
|
prev_node.left = Node(value)
|
|
else:
|
|
prev_node.right = Node(value)
|
|
|
|
def remove(self, value):
|
|
node, prev_node = self.find_existing_node_and_parent(value)
|
|
if prev_node is None and node is not None:
|
|
self.root = None
|
|
return
|
|
elif node is not None:
|
|
if prev_node.value >= value:
|
|
prev_node.left = None
|
|
else:
|
|
prev_node.right = None
|
|
|
|
def search(self, value):
|
|
node = self.root
|
|
|
|
while node is not None:
|
|
if value == node.value:
|
|
return node
|
|
elif value < node.value:
|
|
node = node.left
|
|
else:
|
|
node = node.right
|
|
return None
|
|
|
|
|
|
print("================ Inserts ================")
|
|
|
|
test_tree = BinarySearchTree()
|
|
test_tree.insert(5)
|
|
test_tree.insert(3)
|
|
print(test_tree.root.left.value)
|
|
|
|
test_tree.insert(4)
|
|
test_tree.insert(6)
|
|
print(test_tree.root.left.right.value)
|
|
print(test_tree.root.right.value)
|
|
|
|
print("================ Searches ================")
|
|
|
|
# is the search result giving us the expected node?
|
|
print(test_tree.search(4) == test_tree.root.left.right)
|
|
print(test_tree.search(6) == test_tree.root.right)
|
|
|
|
print("================ Removals ================")
|
|
|
|
test_tree.remove(4)
|
|
print(test_tree.root.left.right)
|
|
test_tree.remove(6)
|
|
print(test_tree.root.right)
|
|
test_tree.remove(3)
|
|
print(test_tree.root.left)
|
|
test_tree.remove(5)
|
|
print(test_tree.root)
|