3 # © Copyright 2021-2022, 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)
117 self, current: Optional[Node], target: Node
118 ) -> List[Optional[Node]]:
119 """Internal helper"""
122 ret: List[Optional[Node]] = [current]
123 if target.value == current.value:
125 elif target.value < current.value:
126 ret.extend(self._parent_path(current.left, target))
129 assert target.value > current.value
130 ret.extend(self._parent_path(current.right, target))
133 def parent_path(self, node: Node) -> List[Optional[Node]]:
134 """Get a node's parent path.
137 node: the node to check
140 a list of nodes representing the path from
141 the tree's root to the node.
145 If the node does not exist in the tree, the last element
146 on the path will be None but the path will indicate the
147 ancestor path of that node were it to be inserted.
149 >>> t = BinarySearchTree()
167 >>> for x in t.parent_path(n):
175 >>> for x in t.parent_path(n):
176 ... if x is not None:
186 return self._parent_path(self.root, node)
188 def __delitem__(self, value: Any) -> bool:
190 Delete an item from the tree and preserve the BST property.
193 value: the value of the node to be deleted.
196 True if the value was found and its associated node was
197 successfully deleted and False otherwise.
199 >>> t = BinarySearchTree()
216 >>> for value in t.iterate_inorder():
226 >>> del t[22] # Note: bool result is discarded
228 >>> for value in t.iterate_inorder():
237 >>> t.__delitem__(13)
239 >>> for value in t.iterate_inorder():
247 >>> t.__delitem__(75)
249 >>> for value in t.iterate_inorder():
261 >>> t.__delitem__(99)
265 if self.root is not None:
266 ret = self._delete(value, None, self.root)
274 def _on_delete(self, parent: Optional[Node], deleted: Node) -> None:
275 """This is called just before deleted is deleted --
276 i.e. before the tree changes."""
279 def _delete(self, value: Any, parent: Optional[Node], node: Node) -> bool:
281 if node.value == value:
283 # Deleting a leaf node
284 if node.left is None and node.right is None:
285 self._on_delete(parent, node)
286 if parent is not None:
287 if parent.left == node:
290 assert parent.right == node
294 # Node only has a right.
295 elif node.left is None:
296 assert node.right is not None
297 self._on_delete(parent, node)
298 if parent is not None:
299 if parent.left == node:
300 parent.left = node.right
302 assert parent.right == node
303 parent.right = node.right
306 # Node only has a left.
307 elif node.right is None:
308 assert node.left is not None
309 self._on_delete(parent, node)
310 if parent is not None:
311 if parent.left == node:
312 parent.left = node.left
314 assert parent.right == node
315 parent.right = node.left
318 # Node has both a left and right.
320 assert node.left is not None and node.right is not None
321 descendent = node.right
322 while descendent.left is not None:
323 descendent = descendent.left
324 node.value = descendent.value
325 return self._delete(node.value, node, node.right)
326 elif value < node.value and node.left is not None:
327 return self._delete(value, node, node.left)
328 elif value > node.value and node.right is not None:
329 return self._delete(value, node, node.right)
335 The count of items in the tree.
337 >>> t = BinarySearchTree()
343 >>> t.__delitem__(50)
359 def __contains__(self, value: Any) -> bool:
362 True if the item is in the tree; False otherwise.
364 return self.__getitem__(value) is not None
366 def _iterate_preorder(self, node: Node):
368 if node.left is not None:
369 yield from self._iterate_preorder(node.left)
370 if node.right is not None:
371 yield from self._iterate_preorder(node.right)
373 def _iterate_inorder(self, node: Node):
374 if node.left is not None:
375 yield from self._iterate_inorder(node.left)
377 if node.right is not None:
378 yield from self._iterate_inorder(node.right)
380 def _iterate_postorder(self, node: Node):
381 if node.left is not None:
382 yield from self._iterate_postorder(node.left)
383 if node.right is not None:
384 yield from self._iterate_postorder(node.right)
387 def iterate_preorder(self):
390 A Generator that yields the tree's items in a
391 preorder traversal sequence.
393 >>> t = BinarySearchTree()
401 >>> for value in t.iterate_preorder():
411 if self.root is not None:
412 yield from self._iterate_preorder(self.root)
414 def iterate_inorder(self):
417 A Generator that yield the tree's items in a preorder
420 >>> t = BinarySearchTree()
437 >>> for value in t.iterate_inorder():
448 if self.root is not None:
449 yield from self._iterate_inorder(self.root)
451 def iterate_postorder(self):
454 A Generator that yield the tree's items in a preorder
457 >>> t = BinarySearchTree()
465 >>> for value in t.iterate_postorder():
475 if self.root is not None:
476 yield from self._iterate_postorder(self.root)
478 def _iterate_leaves(self, node: Node):
479 if node.left is not None:
480 yield from self._iterate_leaves(node.left)
481 if node.right is not None:
482 yield from self._iterate_leaves(node.right)
483 if node.left is None and node.right is None:
486 def iterate_leaves(self):
489 A Gemerator that yielde only the leaf nodes in the
492 >>> t = BinarySearchTree()
500 >>> for value in t.iterate_leaves():
506 if self.root is not None:
507 yield from self._iterate_leaves(self.root)
509 def _iterate_by_depth(self, node: Node, depth: int):
514 if node.left is not None:
515 yield from self._iterate_by_depth(node.left, depth - 1)
516 if node.right is not None:
517 yield from self._iterate_by_depth(node.right, depth - 1)
519 def iterate_nodes_by_depth(self, depth: int) -> Generator[Node, None, None]:
522 depth: the desired depth
525 A Generator that yields nodes at the prescribed depth in
528 >>> t = BinarySearchTree()
536 >>> for value in t.iterate_nodes_by_depth(2):
541 >>> for value in t.iterate_nodes_by_depth(3):
546 if self.root is not None:
547 yield from self._iterate_by_depth(self.root, depth)
549 def get_next_node(self, node: Node) -> Node:
552 node: the node whose next greater successor is desired
555 Given a tree node, returns the next greater node in the tree.
557 >>> t = BinarySearchTree()
575 >>> t.get_next_node(n).value
579 >>> t.get_next_node(n).value
583 if node.right is not None:
585 while x.left is not None:
589 path = self.parent_path(node)
590 assert path[-1] is not None
591 assert path[-1] == node
594 for ancestor in path:
595 assert ancestor is not None
596 if node != ancestor.right:
601 def _depth(self, node: Node, sofar: int) -> int:
602 depth_left = sofar + 1
603 depth_right = sofar + 1
604 if node.left is not None:
605 depth_left = self._depth(node.left, sofar + 1)
606 if node.right is not None:
607 depth_right = self._depth(node.right, sofar + 1)
608 return max(depth_left, depth_right)
610 def depth(self) -> int:
613 The max height (depth) of the tree in plies (edge distance
616 >>> t = BinarySearchTree()
638 if self.root is None:
640 return self._depth(self.root, 0)
642 def height(self) -> int:
643 """Returns the height (i.e. max depth) of the tree"""
650 node: Optional[Node],
651 has_right_sibling: bool,
654 viz = f"\n{padding}{pointer}{node.value}"
655 if has_right_sibling:
660 pointer_right = "└──"
661 if node.right is not None:
666 viz += self.repr_traverse(
667 padding, pointer_left, node.left, node.right is not None
669 viz += self.repr_traverse(padding, pointer_right, node.right, False)
676 An ASCII string representation of the tree.
678 >>> t = BinarySearchTree()
695 if self.root is None:
698 ret = f"{self.root.value}"
699 pointer_right = "└──"
700 if self.root.right is None:
705 ret += self.repr_traverse(
706 "", pointer_left, self.root.left, self.root.left is not None
708 ret += self.repr_traverse("", pointer_right, self.root.right, False)
712 if __name__ == "__main__":