3 from typing import Any, Optional
7 def __init__(self, value: Any) -> None:
13 class BinaryTree(object):
19 def get_root(self) -> Optional[Node]:
22 def insert(self, value: Any):
24 Insert something into the tree.
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.
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 __delitem__(self, value: Any) -> bool:
93 Delete an item from the tree and preserve the BST property.
113 >>> for value in t.iterate_inorder():
123 >>> t.__delitem__(22)
125 >>> for value in t.iterate_inorder():
134 >>> t.__delitem__(13)
136 >>> for value in t.iterate_inorder():
144 >>> t.__delitem__(75)
146 >>> for value in t.iterate_inorder():
153 >>> t.__delitem__(99)
157 if self.root is not None:
158 ret = self._delete(value, None, self.root)
166 def _delete(self, value: Any, parent: Optional[Node], node: Node) -> bool:
168 if node.value == value:
169 # Deleting a leaf node
170 if node.left is None and node.right is None:
171 if parent is not None:
172 if parent.left == node:
175 assert parent.right == node
179 # Node only has a right.
180 elif node.left is None:
181 assert node.right is not None
182 if parent is not None:
183 if parent.left == node:
184 parent.left = node.right
186 assert parent.right == node
187 parent.right = node.right
190 # Node only has a left.
191 elif node.right is None:
192 assert node.left is not None
193 if parent is not None:
194 if parent.left == node:
195 parent.left = node.left
197 assert parent.right == node
198 parent.right = node.left
201 # Node has both a left and right.
203 assert node.left is not None and node.right is not None
204 descendent = node.right
205 while descendent.left is not None:
206 descendent = descendent.left
207 node.value = descendent.value
208 return self._delete(node.value, node, node.right)
209 elif value < node.value and node.left is not None:
210 return self._delete(value, node, node.left)
211 elif value > node.value and node.right is not None:
212 return self._delete(value, node, node.right)
217 Returns the count of items in the tree.
225 >>> t.__delitem__(50)
241 def __contains__(self, value: Any) -> bool:
243 Returns True if the item is in the tree; False otherwise.
246 return self.__getitem__(value) is not None
248 def _iterate_preorder(self, node: Node):
250 if node.left is not None:
251 yield from self._iterate_preorder(node.left)
252 if node.right is not None:
253 yield from self._iterate_preorder(node.right)
255 def _iterate_inorder(self, node: Node):
256 if node.left is not None:
257 yield from self._iterate_inorder(node.left)
259 if node.right is not None:
260 yield from self._iterate_inorder(node.right)
262 def _iterate_postorder(self, node: Node):
263 if node.left is not None:
264 yield from self._iterate_postorder(node.left)
265 if node.right is not None:
266 yield from self._iterate_postorder(node.right)
269 def iterate_preorder(self):
271 Yield the tree's items in a preorder traversal sequence.
281 >>> for value in t.iterate_preorder():
291 if self.root is not None:
292 yield from self._iterate_preorder(self.root)
294 def iterate_inorder(self):
296 Yield the tree's items in a preorder traversal sequence.
306 >>> for value in t.iterate_inorder():
316 if self.root is not None:
317 yield from self._iterate_inorder(self.root)
319 def iterate_postorder(self):
321 Yield the tree's items in a preorder traversal sequence.
331 >>> for value in t.iterate_postorder():
341 if self.root is not None:
342 yield from self._iterate_postorder(self.root)
344 def _iterate_leaves(self, node: Node):
345 if node.left is not None:
346 yield from self._iterate_leaves(node.left)
347 if node.right is not None:
348 yield from self._iterate_leaves(node.right)
349 if node.left is None and node.right is None:
352 def iterate_leaves(self):
354 Iterate only the leaf nodes in the tree.
364 >>> for value in t.iterate_leaves():
370 if self.root is not None:
371 yield from self._iterate_leaves(self.root)
373 def _iterate_by_depth(self, node: Node, depth: int):
378 if node.left is not None:
379 yield from self._iterate_by_depth(node.left, depth - 1)
380 if node.right is not None:
381 yield from self._iterate_by_depth(node.right, depth - 1)
383 def iterate_nodes_by_depth(self, depth: int):
385 Iterate only the leaf nodes in the tree.
395 >>> for value in t.iterate_nodes_by_depth(2):
400 >>> for value in t.iterate_nodes_by_depth(3):
405 if self.root is not None:
406 yield from self._iterate_by_depth(self.root, depth)
408 def _depth(self, node: Node, sofar: int) -> int:
409 depth_left = sofar + 1
410 depth_right = sofar + 1
411 if node.left is not None:
412 depth_left = self._depth(node.left, sofar + 1)
413 if node.right is not None:
414 depth_right = self._depth(node.right, sofar + 1)
415 return max(depth_left, depth_right)
419 Returns the max height (depth) of the tree in plies (edge distance
444 if self.root is None:
446 return self._depth(self.root, 0)
452 if __name__ == '__main__':