3 # © Copyright 2021-2023, Scott Gasch
5 """A binary search tree implementation."""
7 from typing import Generator, List, Optional
9 from pyutils.typez.typing import Comparable
13 def __init__(self, value: Comparable) -> None:
14 """A BST node. Just a left and right reference along with a
15 value. Note that value can be anything as long as it
16 is :class:`Comparable` with other instances of itself.
19 value: a reference to the value of the node. Must be
20 :class:`Comparable` to other values.
23 self.left: Optional[Node] = None
24 self.right: Optional[Node] = None
25 self.value: Comparable = value
28 class BinarySearchTree(object):
34 def get_root(self) -> Optional[Node]:
42 def _on_insert(self, parent: Optional[Node], new: Node) -> None:
43 """This is called immediately _after_ a new node is inserted."""
46 def insert(self, value: Comparable) -> None:
48 Insert something into the tree in :math:`O(log_2 n)` time.
51 value: the value to be inserted.
53 >>> t = BinarySearchTree()
60 >>> t.get_root().value
65 self.root = Node(value)
67 self._on_insert(None, self.root)
69 self._insert(value, self.root)
71 def _insert(self, value: Comparable, node: Node):
72 """Insertion helper"""
73 if value < node.value:
74 if node.left is not None:
75 self._insert(value, node.left)
77 node.left = Node(value)
79 self._on_insert(node, node.left)
81 if node.right is not None:
82 self._insert(value, node.right)
84 node.right = Node(value)
86 self._on_insert(node, node.right)
88 def __getitem__(self, value: Comparable) -> Optional[Node]:
90 Find an item in the tree and return its Node in
91 :math:`O(log_2 n)` time. Returns None if the item is not in
94 >>> t = BinarySearchTree()
106 if self.root is not None:
107 return self._find_exact(value, self.root)
110 def _find_exact(self, target: Comparable, node: Node) -> Optional[Node]:
111 """Recursively traverse the tree looking for a node with the
112 target value. Return that node if it exists, otherwise return
115 if target == node.value:
117 elif target < node.value and node.left is not None:
118 return self._find_exact(target, node.left)
119 elif target > node.value and node.right is not None:
120 return self._find_exact(target, node.right)
123 def _find_lowest_node_less_than_or_equal_to(
124 self, target: Comparable, node: Optional[Node]
126 """Find helper that returns the lowest node that is less
127 than or equal to the target value. Returns None if target is
128 lower than the lowest node in the tree.
130 >>> t = BinarySearchTree()
147 >>> t._find_lowest_node_less_than_or_equal_to(48, t.root).value
149 >>> t._find_lowest_node_less_than_or_equal_to(55, t.root).value
151 >>> t._find_lowest_node_less_than_or_equal_to(100, t.root).value
153 >>> t._find_lowest_node_less_than_or_equal_to(24, t.root).value
155 >>> t._find_lowest_node_less_than_or_equal_to(20, t.root).value
157 >>> t._find_lowest_node_less_than_or_equal_to(72, t.root).value
159 >>> t._find_lowest_node_less_than_or_equal_to(78, t.root).value
161 >>> t._find_lowest_node_less_than_or_equal_to(12, t.root) is None
169 if target == node.value:
172 elif target > node.value:
173 if below := self._find_lowest_node_less_than_or_equal_to(
181 return self._find_lowest_node_less_than_or_equal_to(target, node.left)
183 def _find_lowest_node_greater_than_or_equal_to(
184 self, target: Comparable, node: Optional[Node]
186 """Find helper that returns the lowest node that is greater
187 than or equal to the target value. Returns None if target is
188 higher than the greatest node in the tree.
190 >>> t = BinarySearchTree()
207 >>> t._find_lowest_node_greater_than_or_equal_to(48, t.root).value
209 >>> t._find_lowest_node_greater_than_or_equal_to(55, t.root).value
211 >>> t._find_lowest_node_greater_than_or_equal_to(1, t.root).value
213 >>> t._find_lowest_node_greater_than_or_equal_to(24, t.root).value
215 >>> t._find_lowest_node_greater_than_or_equal_to(20, t.root).value
217 >>> t._find_lowest_node_greater_than_or_equal_to(72, t.root).value
219 >>> t._find_lowest_node_greater_than_or_equal_to(78, t.root).value
221 >>> t._find_lowest_node_greater_than_or_equal_to(95, t.root) is None
229 if target == node.value:
232 elif target > node.value:
233 return self._find_lowest_node_greater_than_or_equal_to(target, node.right)
235 # If target < this node's value, either this node is the
236 # answer or the answer is in this node's left subtree.
238 if below := self._find_lowest_node_greater_than_or_equal_to(
246 self, current: Optional[Node], target: Node
247 ) -> List[Optional[Node]]:
248 """Internal helper"""
251 ret: List[Optional[Node]] = [current]
252 if target.value == current.value:
254 elif target.value < current.value:
255 ret.extend(self._parent_path(current.left, target))
258 assert target.value > current.value
259 ret.extend(self._parent_path(current.right, target))
262 def parent_path(self, node: Node) -> List[Optional[Node]]:
263 """Get a node's parent path in :math:`O(log_2 n)` time.
266 node: the node whose parent path should be returned.
269 a list of nodes representing the path from
270 the tree's root to the given node.
274 If the node does not exist in the tree, the last element
275 on the path will be None but the path will indicate the
276 ancestor path of that node were it to be inserted.
278 >>> t = BinarySearchTree()
296 >>> for x in t.parent_path(n):
304 >>> for x in t.parent_path(n):
305 ... if x is not None:
315 return self._parent_path(self.root, node)
317 def __delitem__(self, value: Comparable) -> bool:
319 Delete an item from the tree and preserve the BST property in
320 :math:`O(log_2 n) time`.
323 value: the value of the node to be deleted.
326 True if the value was found and its associated node was
327 successfully deleted and False otherwise.
329 >>> t = BinarySearchTree()
346 >>> for value in t.iterate_inorder():
356 >>> del t[22] # Note: bool result is discarded
358 >>> for value in t.iterate_inorder():
367 >>> t.__delitem__(13)
369 >>> for value in t.iterate_inorder():
377 >>> t.__delitem__(75)
379 >>> for value in t.iterate_inorder():
391 >>> t.__delitem__(85)
394 >>> t.__delitem__(99)
398 if self.root is not None:
399 ret = self._delete(value, None, self.root)
407 def _on_delete(self, parent: Optional[Node], deleted: Node) -> None:
408 """This is called just after deleted was deleted from the tree"""
411 def _delete(self, value: Comparable, parent: Optional[Node], node: Node) -> bool:
413 if node.value == value:
415 # Deleting a leaf node
416 if node.left is None and node.right is None:
417 if parent is not None:
418 if parent.left == node:
421 assert parent.right == node
423 self._on_delete(parent, node)
426 # Node only has a right.
427 elif node.left is None:
428 assert node.right is not None
429 if parent is not None:
430 if parent.left == node:
431 parent.left = node.right
433 assert parent.right == node
434 parent.right = node.right
435 self._on_delete(parent, node)
438 # Node only has a left.
439 elif node.right is None:
440 assert node.left is not None
441 if parent is not None:
442 if parent.left == node:
443 parent.left = node.left
445 assert parent.right == node
446 parent.right = node.left
447 self._on_delete(parent, node)
450 # Node has both a left and right; get the successor node
451 # to this one and put it here (deleting the successor's
452 # old node). Because these operations are happening only
453 # in the subtree underneath of node, I'm still calling
454 # this delete an O(log_2 n) operation in the docs.
456 assert node.left is not None and node.right is not None
457 successor = self.get_next_node(node)
458 assert successor is not None
459 node.value = successor.value
460 return self._delete(node.value, node, node.right)
462 elif value < node.value and node.left is not None:
463 return self._delete(value, node, node.left)
464 elif value > node.value and node.right is not None:
465 return self._delete(value, node, node.right)
471 The count of items in the tree in :math:`O(1)` time.
473 >>> t = BinarySearchTree()
479 >>> t.__delitem__(50)
495 def __contains__(self, value: Comparable) -> bool:
498 True if the item is in the tree; False otherwise.
500 return self.__getitem__(value) is not None
502 def _iterate_preorder(self, node: Node):
504 if node.left is not None:
505 yield from self._iterate_preorder(node.left)
506 if node.right is not None:
507 yield from self._iterate_preorder(node.right)
509 def _iterate_inorder(self, node: Node):
510 if node.left is not None:
511 yield from self._iterate_inorder(node.left)
513 if node.right is not None:
514 yield from self._iterate_inorder(node.right)
516 def _iterate_postorder(self, node: Node):
517 if node.left is not None:
518 yield from self._iterate_postorder(node.left)
519 if node.right is not None:
520 yield from self._iterate_postorder(node.right)
523 def iterate_preorder(self):
526 A Generator that yields the tree's items in a
527 preorder traversal sequence.
529 >>> t = BinarySearchTree()
537 >>> for value in t.iterate_preorder():
547 if self.root is not None:
548 yield from self._iterate_preorder(self.root)
550 def iterate_inorder(self):
553 A Generator that yield the tree's items in a preorder
556 >>> t = BinarySearchTree()
573 >>> for value in t.iterate_inorder():
584 if self.root is not None:
585 yield from self._iterate_inorder(self.root)
587 def iterate_postorder(self):
590 A Generator that yield the tree's items in a preorder
593 >>> t = BinarySearchTree()
601 >>> for value in t.iterate_postorder():
611 if self.root is not None:
612 yield from self._iterate_postorder(self.root)
614 def _iterate_leaves(self, node: Node):
615 if node.left is not None:
616 yield from self._iterate_leaves(node.left)
617 if node.right is not None:
618 yield from self._iterate_leaves(node.right)
619 if node.left is None and node.right is None:
622 def iterate_leaves(self):
625 A Generator that yields only the leaf nodes in the
628 >>> t = BinarySearchTree()
636 >>> for value in t.iterate_leaves():
642 if self.root is not None:
643 yield from self._iterate_leaves(self.root)
645 def _iterate_by_depth(self, node: Node, depth: int):
650 if node.left is not None:
651 yield from self._iterate_by_depth(node.left, depth - 1)
652 if node.right is not None:
653 yield from self._iterate_by_depth(node.right, depth - 1)
655 def iterate_nodes_by_depth(self, depth: int) -> Generator[Node, None, None]:
658 depth: the desired depth
661 A Generator that yields nodes at the prescribed depth in
664 >>> t = BinarySearchTree()
672 >>> for value in t.iterate_nodes_by_depth(2):
677 >>> for value in t.iterate_nodes_by_depth(3):
682 if self.root is not None:
683 yield from self._iterate_by_depth(self.root, depth)
685 def get_next_node(self, node: Node) -> Optional[Node]:
688 node: the node whose next greater successor is desired
691 Given a tree node, returns the next greater node in the tree.
692 If the given node is the greatest node in the tree, returns None.
694 >>> t = BinarySearchTree()
712 >>> t.get_next_node(n).value
716 >>> t.get_next_node(n).value
720 >>> t.get_next_node(n) is None
724 if node.right is not None:
726 while x.left is not None:
730 path = self.parent_path(node)
731 assert path[-1] is not None
732 assert path[-1] == node
735 for ancestor in path:
736 assert ancestor is not None
737 if node != ancestor.right:
742 def get_nodes_in_range_inclusive(
743 self, lower: Comparable, upper: Comparable
744 ) -> Generator[Node, None, None]:
747 lower: the lower bound of the desired range.
748 upper: the upper bound of the desired range.
751 Generates a sequence of nodes in the desired range.
753 >>> t = BinarySearchTree()
770 >>> for node in t.get_nodes_in_range_inclusive(21, 74):
771 ... print(node.value)
778 node: Optional[Node] = self._find_lowest_node_greater_than_or_equal_to(
782 if lower <= node.value <= upper:
784 node = self.get_next_node(node)
786 def _depth(self, node: Node, sofar: int) -> int:
787 depth_left = sofar + 1
788 depth_right = sofar + 1
789 if node.left is not None:
790 depth_left = self._depth(node.left, sofar + 1)
791 if node.right is not None:
792 depth_right = self._depth(node.right, sofar + 1)
793 return max(depth_left, depth_right)
795 def depth(self) -> int:
798 The max height (depth) of the tree in plies (edge distance
799 from root) in :math:`O(log_2 n)` time.
801 >>> t = BinarySearchTree()
823 if self.root is None:
825 return self._depth(self.root, 0)
827 def height(self) -> int:
828 """Returns the height (i.e. max depth) of the tree"""
835 node: Optional[Node],
836 has_right_sibling: bool,
839 viz = f"\n{padding}{pointer}{node.value}"
840 if has_right_sibling:
845 pointer_right = "└──"
846 if node.right is not None:
851 viz += self.repr_traverse(
852 padding, pointer_left, node.left, node.right is not None
854 viz += self.repr_traverse(padding, pointer_right, node.right, False)
861 An ASCII string representation of the tree.
863 >>> t = BinarySearchTree()
880 if self.root is None:
883 ret = f"{self.root.value}"
884 pointer_right = "└──"
885 if self.root.right is None:
890 ret += self.repr_traverse(
891 "", pointer_left, self.root.left, self.root.left is not None
893 ret += self.repr_traverse("", pointer_right, self.root.right, False)
897 if __name__ == "__main__":