3 from typing import Any, Optional, List
7 def __init__(self, value: Any) -> None:
13 class BinarySearchTree(object):
19 def get_root(self) -> Optional[Node]:
22 def insert(self, value: Any):
24 Insert something into the tree.
26 >>> t = BinarySearchTree()
33 >>> t.get_root().value
38 self.root = Node(value)
41 self._insert(value, self.root)
43 def _insert(self, value: Any, node: Node):
44 """Insertion helper"""
45 if value < node.value:
46 if node.left is not None:
47 self._insert(value, node.left)
49 node.left = Node(value)
52 if node.right is not None:
53 self._insert(value, node.right)
55 node.right = Node(value)
58 def __getitem__(self, value: Any) -> Optional[Node]:
60 Find an item in the tree and return its Node. Returns
61 None if the item is not in the tree.
63 >>> t = BinarySearchTree()
75 if self.root is not None:
76 return self._find(value, self.root)
79 def _find(self, value: Any, node: Node) -> Optional[Node]:
81 if value == node.value:
83 elif (value < node.value and node.left is not None):
84 return self._find(value, node.left)
86 assert value > node.value
87 if node.right is not None:
88 return self._find(value, node.right)
91 def _parent_path(self, current: Node, target: Node):
95 if target.value == current.value:
97 elif target.value < current.value:
98 ret.extend(self._parent_path(current.left, target))
101 assert target.value > current.value
102 ret.extend(self._parent_path(current.right, target))
105 def parent_path(self, node: Node) -> Optional[List[Node]]:
106 """Return a list of nodes representing the path from
107 the tree's root to the node argument.
109 >>> t = BinarySearchTree()
127 >>> for x in t.parent_path(n):
135 return self._parent_path(self.root, node)
137 def __delitem__(self, value: Any) -> bool:
139 Delete an item from the tree and preserve the BST property.
150 >>> t = BinarySearchTree()
159 >>> for value in t.iterate_inorder():
169 >>> del t[22] # Note: bool result is discarded
171 >>> for value in t.iterate_inorder():
180 >>> t.__delitem__(13)
182 >>> for value in t.iterate_inorder():
190 >>> t.__delitem__(75)
192 >>> for value in t.iterate_inorder():
199 >>> t.__delitem__(99)
203 if self.root is not None:
204 ret = self._delete(value, None, self.root)
212 def _delete(self, value: Any, parent: Optional[Node], node: Node) -> bool:
214 if node.value == value:
215 # Deleting a leaf node
216 if node.left is None and node.right is None:
217 if parent is not None:
218 if parent.left == node:
221 assert parent.right == node
225 # Node only has a right.
226 elif node.left is None:
227 assert node.right is not None
228 if parent is not None:
229 if parent.left == node:
230 parent.left = node.right
232 assert parent.right == node
233 parent.right = node.right
236 # Node only has a left.
237 elif node.right is None:
238 assert node.left is not None
239 if parent is not None:
240 if parent.left == node:
241 parent.left = node.left
243 assert parent.right == node
244 parent.right = node.left
247 # Node has both a left and right.
249 assert node.left is not None and node.right is not None
250 descendent = node.right
251 while descendent.left is not None:
252 descendent = descendent.left
253 node.value = descendent.value
254 return self._delete(node.value, node, node.right)
255 elif value < node.value and node.left is not None:
256 return self._delete(value, node, node.left)
257 elif value > node.value and node.right is not None:
258 return self._delete(value, node, node.right)
263 Returns the count of items in the tree.
265 >>> t = BinarySearchTree()
271 >>> t.__delitem__(50)
287 def __contains__(self, value: Any) -> bool:
289 Returns True if the item is in the tree; False otherwise.
292 return self.__getitem__(value) is not None
294 def _iterate_preorder(self, node: Node):
296 if node.left is not None:
297 yield from self._iterate_preorder(node.left)
298 if node.right is not None:
299 yield from self._iterate_preorder(node.right)
301 def _iterate_inorder(self, node: Node):
302 if node.left is not None:
303 yield from self._iterate_inorder(node.left)
305 if node.right is not None:
306 yield from self._iterate_inorder(node.right)
308 def _iterate_postorder(self, node: Node):
309 if node.left is not None:
310 yield from self._iterate_postorder(node.left)
311 if node.right is not None:
312 yield from self._iterate_postorder(node.right)
315 def iterate_preorder(self):
317 Yield the tree's items in a preorder traversal sequence.
319 >>> t = BinarySearchTree()
327 >>> for value in t.iterate_preorder():
337 if self.root is not None:
338 yield from self._iterate_preorder(self.root)
340 def iterate_inorder(self):
342 Yield the tree's items in a preorder traversal sequence.
344 >>> t = BinarySearchTree()
361 >>> for value in t.iterate_inorder():
372 if self.root is not None:
373 yield from self._iterate_inorder(self.root)
375 def iterate_postorder(self):
377 Yield the tree's items in a preorder traversal sequence.
379 >>> t = BinarySearchTree()
387 >>> for value in t.iterate_postorder():
397 if self.root is not None:
398 yield from self._iterate_postorder(self.root)
400 def _iterate_leaves(self, node: Node):
401 if node.left is not None:
402 yield from self._iterate_leaves(node.left)
403 if node.right is not None:
404 yield from self._iterate_leaves(node.right)
405 if node.left is None and node.right is None:
408 def iterate_leaves(self):
410 Iterate only the leaf nodes in the tree.
412 >>> t = BinarySearchTree()
420 >>> for value in t.iterate_leaves():
426 if self.root is not None:
427 yield from self._iterate_leaves(self.root)
429 def _iterate_by_depth(self, node: Node, depth: int):
434 if node.left is not None:
435 yield from self._iterate_by_depth(node.left, depth - 1)
436 if node.right is not None:
437 yield from self._iterate_by_depth(node.right, depth - 1)
439 def iterate_nodes_by_depth(self, depth: int):
441 Iterate only the leaf nodes in the tree.
443 >>> t = BinarySearchTree()
451 >>> for value in t.iterate_nodes_by_depth(2):
456 >>> for value in t.iterate_nodes_by_depth(3):
461 if self.root is not None:
462 yield from self._iterate_by_depth(self.root, depth)
464 def get_next_node(self, node: Node) -> Node:
466 Given a tree node, get the next greater node in the tree.
468 >>> t = BinarySearchTree()
486 >>> t.get_next_node(n).value
490 >>> t.get_next_node(n).value
494 if node.right is not None:
496 while x.left is not None:
500 path = self.parent_path(node)
501 assert path[-1] == node
504 for ancestor in path:
505 if node != ancestor.right:
509 def _depth(self, node: Node, sofar: int) -> int:
510 depth_left = sofar + 1
511 depth_right = sofar + 1
512 if node.left is not None:
513 depth_left = self._depth(node.left, sofar + 1)
514 if node.right is not None:
515 depth_right = self._depth(node.right, sofar + 1)
516 return max(depth_left, depth_right)
520 Returns the max height (depth) of the tree in plies (edge distance
523 >>> t = BinarySearchTree()
545 if self.root is None:
547 return self._depth(self.root, 0)
552 def repr_traverse(self, padding: str, pointer: str, node: Node, has_right_sibling: bool) -> str:
554 self.viz += f'\n{padding}{pointer}{node.value}'
555 if has_right_sibling:
560 pointer_right = "└──"
561 if node.right is not None:
566 self.repr_traverse(padding, pointer_left, node.left, node.right is not None)
567 self.repr_traverse(padding, pointer_right, node.right, False)
571 Draw the tree in ASCII.
573 >>> t = BinarySearchTree()
590 if self.root is None:
593 self.viz = f'{self.root.value}'
594 pointer_right = "└──"
595 if self.root.right is None:
600 self.repr_traverse('', pointer_left, self.root.left, self.root.left is not None)
601 self.repr_traverse('', pointer_right, self.root.right, False)
605 if __name__ == '__main__':