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 insert(self, value: Any) -> None:
41 Insert something into the tree.
44 value: the value to be inserted.
46 >>> t = BinarySearchTree()
53 >>> t.get_root().value
58 self.root = Node(value)
61 self._insert(value, self.root)
63 def _insert(self, value: Any, node: Node):
64 """Insertion helper"""
65 if value < node.value:
66 if node.left is not None:
67 self._insert(value, node.left)
69 node.left = Node(value)
72 if node.right is not None:
73 self._insert(value, node.right)
75 node.right = Node(value)
78 def __getitem__(self, value: Any) -> Optional[Node]:
80 Find an item in the tree and return its Node. Returns
81 None if the item is not in the tree.
83 >>> t = BinarySearchTree()
95 if self.root is not None:
96 return self._find(value, self.root)
99 def _find(self, value: Any, node: Node) -> Optional[Node]:
101 if value == node.value:
103 elif value < node.value and node.left is not None:
104 return self._find(value, node.left)
105 elif value > node.value and node.right is not None:
106 return self._find(value, node.right)
110 self, current: Optional[Node], target: Node
111 ) -> List[Optional[Node]]:
112 """Internal helper"""
115 ret: List[Optional[Node]] = [current]
116 if target.value == current.value:
118 elif target.value < current.value:
119 ret.extend(self._parent_path(current.left, target))
122 assert target.value > current.value
123 ret.extend(self._parent_path(current.right, target))
126 def parent_path(self, node: Node) -> List[Optional[Node]]:
127 """Get a node's parent path.
130 node: the node to check
133 a list of nodes representing the path from
134 the tree's root to the node.
138 If the node does not exist in the tree, the last element
139 on the path will be None but the path will indicate the
140 ancestor path of that node were it to be inserted.
142 >>> t = BinarySearchTree()
160 >>> for x in t.parent_path(n):
168 >>> for x in t.parent_path(n):
169 ... if x is not None:
179 return self._parent_path(self.root, node)
181 def __delitem__(self, value: Any) -> bool:
183 Delete an item from the tree and preserve the BST property.
186 value: the value of the node to be deleted.
189 True if the value was found and its associated node was
190 successfully deleted and False otherwise.
192 >>> t = BinarySearchTree()
209 >>> for value in t.iterate_inorder():
219 >>> del t[22] # Note: bool result is discarded
221 >>> for value in t.iterate_inorder():
230 >>> t.__delitem__(13)
232 >>> for value in t.iterate_inorder():
240 >>> t.__delitem__(75)
242 >>> for value in t.iterate_inorder():
254 >>> t.__delitem__(99)
258 if self.root is not None:
259 ret = self._delete(value, None, self.root)
267 def _delete(self, value: Any, parent: Optional[Node], node: Node) -> bool:
269 if node.value == value:
270 # Deleting a leaf node
271 if node.left is None and node.right is None:
272 if parent is not None:
273 if parent.left == node:
276 assert parent.right == node
280 # Node only has a right.
281 elif node.left is None:
282 assert node.right is not None
283 if parent is not None:
284 if parent.left == node:
285 parent.left = node.right
287 assert parent.right == node
288 parent.right = node.right
291 # Node only has a left.
292 elif node.right is None:
293 assert node.left is not None
294 if parent is not None:
295 if parent.left == node:
296 parent.left = node.left
298 assert parent.right == node
299 parent.right = node.left
302 # Node has both a left and right.
304 assert node.left is not None and node.right is not None
305 descendent = node.right
306 while descendent.left is not None:
307 descendent = descendent.left
308 node.value = descendent.value
309 return self._delete(node.value, node, node.right)
310 elif value < node.value and node.left is not None:
311 return self._delete(value, node, node.left)
312 elif value > node.value and node.right is not None:
313 return self._delete(value, node, node.right)
319 The count of items in the tree.
321 >>> t = BinarySearchTree()
327 >>> t.__delitem__(50)
343 def __contains__(self, value: Any) -> bool:
346 True if the item is in the tree; False otherwise.
348 return self.__getitem__(value) is not None
350 def _iterate_preorder(self, node: Node):
352 if node.left is not None:
353 yield from self._iterate_preorder(node.left)
354 if node.right is not None:
355 yield from self._iterate_preorder(node.right)
357 def _iterate_inorder(self, node: Node):
358 if node.left is not None:
359 yield from self._iterate_inorder(node.left)
361 if node.right is not None:
362 yield from self._iterate_inorder(node.right)
364 def _iterate_postorder(self, node: Node):
365 if node.left is not None:
366 yield from self._iterate_postorder(node.left)
367 if node.right is not None:
368 yield from self._iterate_postorder(node.right)
371 def iterate_preorder(self):
374 A Generator that yields the tree's items in a
375 preorder traversal sequence.
377 >>> t = BinarySearchTree()
385 >>> for value in t.iterate_preorder():
395 if self.root is not None:
396 yield from self._iterate_preorder(self.root)
398 def iterate_inorder(self):
401 A Generator that yield the tree's items in a preorder
404 >>> t = BinarySearchTree()
421 >>> for value in t.iterate_inorder():
432 if self.root is not None:
433 yield from self._iterate_inorder(self.root)
435 def iterate_postorder(self):
438 A Generator that yield the tree's items in a preorder
441 >>> t = BinarySearchTree()
449 >>> for value in t.iterate_postorder():
459 if self.root is not None:
460 yield from self._iterate_postorder(self.root)
462 def _iterate_leaves(self, node: Node):
463 if node.left is not None:
464 yield from self._iterate_leaves(node.left)
465 if node.right is not None:
466 yield from self._iterate_leaves(node.right)
467 if node.left is None and node.right is None:
470 def iterate_leaves(self):
473 A Gemerator that yielde only the leaf nodes in the
476 >>> t = BinarySearchTree()
484 >>> for value in t.iterate_leaves():
490 if self.root is not None:
491 yield from self._iterate_leaves(self.root)
493 def _iterate_by_depth(self, node: Node, depth: int):
498 if node.left is not None:
499 yield from self._iterate_by_depth(node.left, depth - 1)
500 if node.right is not None:
501 yield from self._iterate_by_depth(node.right, depth - 1)
503 def iterate_nodes_by_depth(self, depth: int) -> Generator[Node, None, None]:
506 depth: the desired depth
509 A Generator that yields nodes at the prescribed depth in
512 >>> t = BinarySearchTree()
520 >>> for value in t.iterate_nodes_by_depth(2):
525 >>> for value in t.iterate_nodes_by_depth(3):
530 if self.root is not None:
531 yield from self._iterate_by_depth(self.root, depth)
533 def get_next_node(self, node: Node) -> Node:
536 node: the node whose next greater successor is desired
539 Given a tree node, returns the next greater node in the tree.
541 >>> t = BinarySearchTree()
559 >>> t.get_next_node(n).value
563 >>> t.get_next_node(n).value
567 if node.right is not None:
569 while x.left is not None:
573 path = self.parent_path(node)
574 assert path[-1] is not None
575 assert path[-1] == node
578 for ancestor in path:
579 assert ancestor is not None
580 if node != ancestor.right:
585 def _depth(self, node: Node, sofar: int) -> int:
586 depth_left = sofar + 1
587 depth_right = sofar + 1
588 if node.left is not None:
589 depth_left = self._depth(node.left, sofar + 1)
590 if node.right is not None:
591 depth_right = self._depth(node.right, sofar + 1)
592 return max(depth_left, depth_right)
594 def depth(self) -> int:
597 The max height (depth) of the tree in plies (edge distance
600 >>> t = BinarySearchTree()
622 if self.root is None:
624 return self._depth(self.root, 0)
626 def height(self) -> int:
627 """Returns the height (i.e. max depth) of the tree"""
634 node: Optional[Node],
635 has_right_sibling: bool,
638 viz = f'\n{padding}{pointer}{node.value}'
639 if has_right_sibling:
644 pointer_right = "└──"
645 if node.right is not None:
650 viz += self.repr_traverse(
651 padding, pointer_left, node.left, node.right is not None
653 viz += self.repr_traverse(padding, pointer_right, node.right, False)
660 An ASCII string representation of the tree.
662 >>> t = BinarySearchTree()
679 if self.root is None:
682 ret = f'{self.root.value}'
683 pointer_right = "└──"
684 if self.root.right is None:
689 ret += self.repr_traverse(
690 '', pointer_left, self.root.left, self.root.left is not None
692 ret += self.repr_traverse('', pointer_right, self.root.right, False)
696 if __name__ == '__main__':