3 # © Copyright 2021-2023, Scott Gasch
5 """A binary search tree implementation."""
7 from typing import Any, Generator, List, Optional
11 def __init__(self, value: Any) -> None:
13 A BST node. Note that value can be anything as long as it
14 is comparable. Check out :meth:`functools.total_ordering`
15 (https://docs.python.org/3/library/functools.html#functools.total_ordering)
18 value: a reference to the value of the node.
20 self.left: Optional[Node] = None
21 self.right: Optional[Node] = None
25 class BinarySearchTree(object):
31 def get_root(self) -> Optional[Node]:
39 def _on_insert(self, parent: Optional[Node], new: Node) -> None:
40 """This is called immediately _after_ a new node is inserted."""
43 def insert(self, value: Any) -> None:
45 Insert something into the tree.
48 value: the value to be inserted.
50 >>> t = BinarySearchTree()
57 >>> t.get_root().value
62 self.root = Node(value)
64 self._on_insert(None, self.root)
66 self._insert(value, self.root)
68 def _insert(self, value: Any, node: Node):
69 """Insertion helper"""
70 if value < node.value:
71 if node.left is not None:
72 self._insert(value, node.left)
74 node.left = Node(value)
76 self._on_insert(node, node.left)
78 if node.right is not None:
79 self._insert(value, node.right)
81 node.right = Node(value)
83 self._on_insert(node, node.right)
85 def __getitem__(self, value: Any) -> Optional[Node]:
87 Find an item in the tree and return its Node. Returns
88 None if the item is not in the tree.
90 >>> t = BinarySearchTree()
102 if self.root is not None:
103 return self._find(value, self.root)
106 def _find(self, value: Any, node: Node) -> Optional[Node]:
108 if value == node.value:
110 elif value < node.value and node.left is not None:
111 return self._find(value, node.left)
112 elif value > node.value and node.right is not None:
113 return self._find(value, node.right)
116 def _find_fuzzy(self, value: Any, node: Node) -> Node:
117 """Find helper that returns the closest node <= the target value."""
118 if value == node.value:
121 if value < node.value and node.left is not None:
122 below = self._find_fuzzy(value, node.left)
123 elif value > node.value and node.right is not None:
124 below = self._find_fuzzy(value, node.right)
130 self, current: Optional[Node], target: Node
131 ) -> List[Optional[Node]]:
132 """Internal helper"""
135 ret: List[Optional[Node]] = [current]
136 if target.value == current.value:
138 elif target.value < current.value:
139 ret.extend(self._parent_path(current.left, target))
142 assert target.value > current.value
143 ret.extend(self._parent_path(current.right, target))
146 def parent_path(self, node: Node) -> List[Optional[Node]]:
147 """Get a node's parent path.
150 node: the node to check
153 a list of nodes representing the path from
154 the tree's root to the node.
158 If the node does not exist in the tree, the last element
159 on the path will be None but the path will indicate the
160 ancestor path of that node were it to be inserted.
162 >>> t = BinarySearchTree()
180 >>> for x in t.parent_path(n):
188 >>> for x in t.parent_path(n):
189 ... if x is not None:
199 return self._parent_path(self.root, node)
201 def __delitem__(self, value: Any) -> bool:
203 Delete an item from the tree and preserve the BST property.
206 value: the value of the node to be deleted.
209 True if the value was found and its associated node was
210 successfully deleted and False otherwise.
212 >>> t = BinarySearchTree()
229 >>> for value in t.iterate_inorder():
239 >>> del t[22] # Note: bool result is discarded
241 >>> for value in t.iterate_inorder():
250 >>> t.__delitem__(13)
252 >>> for value in t.iterate_inorder():
260 >>> t.__delitem__(75)
262 >>> for value in t.iterate_inorder():
274 >>> t.__delitem__(99)
278 if self.root is not None:
279 ret = self._delete(value, None, self.root)
287 def _on_delete(self, parent: Optional[Node], deleted: Node) -> None:
288 """This is called just after deleted was deleted from the tree"""
291 def _delete(self, value: Any, parent: Optional[Node], node: Node) -> bool:
293 if node.value == value:
295 # Deleting a leaf node
296 if node.left is None and node.right is None:
297 if parent is not None:
298 if parent.left == node:
301 assert parent.right == node
303 self._on_delete(parent, node)
306 # Node only has a right.
307 elif node.left is None:
308 assert node.right is not None
309 if parent is not None:
310 if parent.left == node:
311 parent.left = node.right
313 assert parent.right == node
314 parent.right = node.right
315 self._on_delete(parent, node)
318 # Node only has a left.
319 elif node.right is None:
320 assert node.left is not None
321 if parent is not None:
322 if parent.left == node:
323 parent.left = node.left
325 assert parent.right == node
326 parent.right = node.left
327 self._on_delete(parent, node)
330 # Node has both a left and right.
332 assert node.left is not None and node.right is not None
333 descendent = node.right
334 while descendent.left is not None:
335 descendent = descendent.left
336 node.value = descendent.value
337 return self._delete(node.value, node, node.right)
338 elif value < node.value and node.left is not None:
339 return self._delete(value, node, node.left)
340 elif value > node.value and node.right is not None:
341 return self._delete(value, node, node.right)
347 The count of items in the tree.
349 >>> t = BinarySearchTree()
355 >>> t.__delitem__(50)
371 def __contains__(self, value: Any) -> bool:
374 True if the item is in the tree; False otherwise.
376 return self.__getitem__(value) is not None
378 def _iterate_preorder(self, node: Node):
380 if node.left is not None:
381 yield from self._iterate_preorder(node.left)
382 if node.right is not None:
383 yield from self._iterate_preorder(node.right)
385 def _iterate_inorder(self, node: Node):
386 if node.left is not None:
387 yield from self._iterate_inorder(node.left)
389 if node.right is not None:
390 yield from self._iterate_inorder(node.right)
392 def _iterate_postorder(self, node: Node):
393 if node.left is not None:
394 yield from self._iterate_postorder(node.left)
395 if node.right is not None:
396 yield from self._iterate_postorder(node.right)
399 def iterate_preorder(self):
402 A Generator that yields the tree's items in a
403 preorder traversal sequence.
405 >>> t = BinarySearchTree()
413 >>> for value in t.iterate_preorder():
423 if self.root is not None:
424 yield from self._iterate_preorder(self.root)
426 def iterate_inorder(self):
429 A Generator that yield the tree's items in a preorder
432 >>> t = BinarySearchTree()
449 >>> for value in t.iterate_inorder():
460 if self.root is not None:
461 yield from self._iterate_inorder(self.root)
463 def iterate_postorder(self):
466 A Generator that yield the tree's items in a preorder
469 >>> t = BinarySearchTree()
477 >>> for value in t.iterate_postorder():
487 if self.root is not None:
488 yield from self._iterate_postorder(self.root)
490 def _iterate_leaves(self, node: Node):
491 if node.left is not None:
492 yield from self._iterate_leaves(node.left)
493 if node.right is not None:
494 yield from self._iterate_leaves(node.right)
495 if node.left is None and node.right is None:
498 def iterate_leaves(self):
501 A Gemerator that yielde only the leaf nodes in the
504 >>> t = BinarySearchTree()
512 >>> for value in t.iterate_leaves():
518 if self.root is not None:
519 yield from self._iterate_leaves(self.root)
521 def _iterate_by_depth(self, node: Node, depth: int):
526 if node.left is not None:
527 yield from self._iterate_by_depth(node.left, depth - 1)
528 if node.right is not None:
529 yield from self._iterate_by_depth(node.right, depth - 1)
531 def iterate_nodes_by_depth(self, depth: int) -> Generator[Node, None, None]:
534 depth: the desired depth
537 A Generator that yields nodes at the prescribed depth in
540 >>> t = BinarySearchTree()
548 >>> for value in t.iterate_nodes_by_depth(2):
553 >>> for value in t.iterate_nodes_by_depth(3):
558 if self.root is not None:
559 yield from self._iterate_by_depth(self.root, depth)
561 def get_next_node(self, node: Node) -> Optional[Node]:
564 node: the node whose next greater successor is desired
567 Given a tree node, returns the next greater node in the tree.
568 If the given node is the greatest node in the tree, returns None.
570 >>> t = BinarySearchTree()
588 >>> t.get_next_node(n).value
592 >>> t.get_next_node(n).value
596 >>> t.get_next_node(n) is None
600 if node.right is not None:
602 while x.left is not None:
606 path = self.parent_path(node)
607 assert path[-1] is not None
608 assert path[-1] == node
611 for ancestor in path:
612 assert ancestor is not None
613 if node != ancestor.right:
618 def get_nodes_in_range_inclusive(self, lower: Any, upper: Any):
620 >>> t = BinarySearchTree()
637 >>> for node in t.get_nodes_in_range_inclusive(21, 74):
638 ... print(node.value)
645 node: Optional[Node] = self._find_fuzzy(lower, self.root)
647 if lower <= node.value <= upper:
649 node = self.get_next_node(node)
651 def _depth(self, node: Node, sofar: int) -> int:
652 depth_left = sofar + 1
653 depth_right = sofar + 1
654 if node.left is not None:
655 depth_left = self._depth(node.left, sofar + 1)
656 if node.right is not None:
657 depth_right = self._depth(node.right, sofar + 1)
658 return max(depth_left, depth_right)
660 def depth(self) -> int:
663 The max height (depth) of the tree in plies (edge distance
666 >>> t = BinarySearchTree()
688 if self.root is None:
690 return self._depth(self.root, 0)
692 def height(self) -> int:
693 """Returns the height (i.e. max depth) of the tree"""
700 node: Optional[Node],
701 has_right_sibling: bool,
704 viz = f"\n{padding}{pointer}{node.value}"
705 if has_right_sibling:
710 pointer_right = "└──"
711 if node.right is not None:
716 viz += self.repr_traverse(
717 padding, pointer_left, node.left, node.right is not None
719 viz += self.repr_traverse(padding, pointer_right, node.right, False)
726 An ASCII string representation of the tree.
728 >>> t = BinarySearchTree()
745 if self.root is None:
748 ret = f"{self.root.value}"
749 pointer_right = "└──"
750 if self.root.right is None:
755 ret += self.repr_traverse(
756 "", pointer_left, self.root.left, self.root.left is not None
758 ret += self.repr_traverse("", pointer_right, self.root.right, False)
762 if __name__ == "__main__":