--- /dev/null
+#!/usr/bin/env python3
+
+from typing import Any, Optional
+
+
+class Node(object):
+ def __init__(self, value: Any) -> None:
+ self.left = None
+ self.right = None
+ self.value = value
+
+
+class BinaryTree(object):
+ def __init__(self):
+ self.root = None
+ self.count = 0
+ self.traverse = None
+
+ def get_root(self) -> Optional[Node]:
+ return self.root
+
+ def insert(self, value: Any):
+ """
+ Insert something into the tree.
+
+ >>> t = BinaryTree()
+ >>> t.insert(10)
+ >>> t.insert(20)
+ >>> t.insert(5)
+ >>> len(t)
+ 3
+
+ >>> t.get_root().value
+ 10
+
+ """
+ if self.root is None:
+ self.root = Node(value)
+ self.count = 1
+ else:
+ self._insert(value, self.root)
+
+ def _insert(self, value: Any, node: Node):
+ """Insertion helper"""
+ if value < node.value:
+ if node.left is not None:
+ self._insert(value, node.left)
+ else:
+ node.left = Node(value)
+ self.count += 1
+ else:
+ if node.right is not None:
+ self._insert(value, node.right)
+ else:
+ node.right = Node(value)
+ self.count += 1
+
+ def __getitem__(self, value: Any) -> Optional[Node]:
+ """
+ Find an item in the tree and return its Node. Returns
+ None if the item is not in the tree.
+
+ >>> t = BinaryTree()
+ >>> t[99]
+
+ >>> t.insert(10)
+ >>> t.insert(20)
+ >>> t.insert(5)
+ >>> t[10].value
+ 10
+
+ >>> t[99]
+
+ """
+ if self.root is not None:
+ return self._find(value, self.root)
+ return None
+
+ def _find(self, value: Any, node: Node) -> Optional[Node]:
+ """Find helper"""
+ if value == node.value:
+ return node
+ elif (value < node.value and node.left is not None):
+ return self._find(value, node.left)
+ else:
+ assert value > node.value
+ if node.right is not None:
+ return self._find(value, node.right)
+ return None
+
+ def __delitem__(self, value: Any) -> bool:
+ """
+ Delete an item from the tree and preserve the BST property.
+
+ 50
+ / \
+ 25 75
+ / / \
+ 22 66 85
+ /
+ 13
+
+
+ >>> t = BinaryTree()
+ >>> t.insert(50)
+ >>> t.insert(75)
+ >>> t.insert(25)
+ >>> t.insert(66)
+ >>> t.insert(22)
+ >>> t.insert(13)
+ >>> t.insert(85)
+
+ >>> for value in t.iterate_inorder():
+ ... print(value)
+ 13
+ 22
+ 25
+ 50
+ 66
+ 75
+ 85
+
+ >>> t.__delitem__(22)
+ True
+ >>> for value in t.iterate_inorder():
+ ... print(value)
+ 13
+ 25
+ 50
+ 66
+ 75
+ 85
+
+ >>> t.__delitem__(13)
+ True
+ >>> for value in t.iterate_inorder():
+ ... print(value)
+ 25
+ 50
+ 66
+ 75
+ 85
+
+ >>> t.__delitem__(75)
+ True
+ >>> for value in t.iterate_inorder():
+ ... print(value)
+ 25
+ 50
+ 66
+ 85
+
+ >>> t.__delitem__(99)
+ False
+
+ """
+ if self.root is not None:
+ ret = self._delete(value, None, self.root)
+ if ret:
+ self.count -= 1
+ if self.count == 0:
+ self.root = None
+ return ret
+ return False
+
+ def _delete(self, value: Any, parent: Optional[Node], node: Node) -> bool:
+ """Delete helper"""
+ if node.value == value:
+ # Deleting a leaf node
+ if node.left is None and node.right is None:
+ if parent is not None:
+ if parent.left == node:
+ parent.left = None
+ else:
+ assert parent.right == node
+ parent.right = None
+ return True
+
+ # Node only has a right.
+ elif node.left is None:
+ assert node.right is not None
+ if parent is not None:
+ if parent.left == node:
+ parent.left = node.right
+ else:
+ assert parent.right == node
+ parent.right = node.right
+ return True
+
+ # Node only has a left.
+ elif node.right is None:
+ assert node.left is not None
+ if parent is not None:
+ if parent.left == node:
+ parent.left = node.left
+ else:
+ assert parent.right == node
+ parent.right = node.left
+ return True
+
+ # Node has both a left and right.
+ else:
+ assert node.left is not None and node.right is not None
+ descendent = node.right
+ while descendent.left is not None:
+ descendent = descendent.left
+ node.value = descendent.value
+ return self._delete(node.value, node, node.right)
+ elif value < node.value and node.left is not None:
+ return self._delete(value, node, node.left)
+ elif value > node.value and node.right is not None:
+ return self._delete(value, node, node.right)
+ return False
+
+ def __len__(self):
+ """
+ Returns the count of items in the tree.
+
+ >>> t = BinaryTree()
+ >>> len(t)
+ 0
+ >>> t.insert(50)
+ >>> len(t)
+ 1
+ >>> t.__delitem__(50)
+ True
+ >>> len(t)
+ 0
+ >>> t.insert(75)
+ >>> t.insert(25)
+ >>> t.insert(66)
+ >>> t.insert(22)
+ >>> t.insert(13)
+ >>> t.insert(85)
+ >>> len(t)
+ 6
+
+ """
+ return self.count
+
+ def __contains__(self, value: Any) -> bool:
+ """
+ Returns True if the item is in the tree; False otherwise.
+
+ """
+ return self.__getitem__(value) is not None
+
+ def _iterate_preorder(self, node: Node):
+ yield node.value
+ if node.left is not None:
+ yield from self._iterate_preorder(node.left)
+ if node.right is not None:
+ yield from self._iterate_preorder(node.right)
+
+ def _iterate_inorder(self, node: Node):
+ if node.left is not None:
+ yield from self._iterate_inorder(node.left)
+ yield node.value
+ if node.right is not None:
+ yield from self._iterate_inorder(node.right)
+
+ def _iterate_postorder(self, node: Node):
+ if node.left is not None:
+ yield from self._iterate_postorder(node.left)
+ if node.right is not None:
+ yield from self._iterate_postorder(node.right)
+ yield node.value
+
+ def iterate_preorder(self):
+ """
+ Yield the tree's items in a preorder traversal sequence.
+
+ >>> t = BinaryTree()
+ >>> t.insert(50)
+ >>> t.insert(75)
+ >>> t.insert(25)
+ >>> t.insert(66)
+ >>> t.insert(22)
+ >>> t.insert(13)
+
+ >>> for value in t.iterate_preorder():
+ ... print(value)
+ 50
+ 25
+ 22
+ 13
+ 75
+ 66
+
+ """
+ if self.root is not None:
+ yield from self._iterate_preorder(self.root)
+
+ def iterate_inorder(self):
+ """
+ Yield the tree's items in a preorder traversal sequence.
+
+ >>> t = BinaryTree()
+ >>> t.insert(50)
+ >>> t.insert(75)
+ >>> t.insert(25)
+ >>> t.insert(66)
+ >>> t.insert(22)
+ >>> t.insert(13)
+
+ >>> for value in t.iterate_inorder():
+ ... print(value)
+ 13
+ 22
+ 25
+ 50
+ 66
+ 75
+
+ """
+ if self.root is not None:
+ yield from self._iterate_inorder(self.root)
+
+ def iterate_postorder(self):
+ """
+ Yield the tree's items in a preorder traversal sequence.
+
+ >>> t = BinaryTree()
+ >>> t.insert(50)
+ >>> t.insert(75)
+ >>> t.insert(25)
+ >>> t.insert(66)
+ >>> t.insert(22)
+ >>> t.insert(13)
+
+ >>> for value in t.iterate_postorder():
+ ... print(value)
+ 13
+ 22
+ 25
+ 66
+ 75
+ 50
+
+ """
+ if self.root is not None:
+ yield from self._iterate_postorder(self.root)
+
+ def _iterate_leaves(self, node: Node):
+ if node.left is not None:
+ yield from self._iterate_leaves(node.left)
+ if node.right is not None:
+ yield from self._iterate_leaves(node.right)
+ if node.left is None and node.right is None:
+ yield node.value
+
+ def iterate_leaves(self):
+ """
+ Iterate only the leaf nodes in the tree.
+
+ >>> t = BinaryTree()
+ >>> t.insert(50)
+ >>> t.insert(75)
+ >>> t.insert(25)
+ >>> t.insert(66)
+ >>> t.insert(22)
+ >>> t.insert(13)
+
+ >>> for value in t.iterate_leaves():
+ ... print(value)
+ 13
+ 66
+
+ """
+ if self.root is not None:
+ yield from self._iterate_leaves(self.root)
+
+ def _iterate_by_depth(self, node: Node, depth: int):
+ if depth == 0:
+ yield node.value
+ else:
+ assert depth > 0
+ if node.left is not None:
+ yield from self._iterate_by_depth(node.left, depth - 1)
+ if node.right is not None:
+ yield from self._iterate_by_depth(node.right, depth - 1)
+
+ def iterate_nodes_by_depth(self, depth: int):
+ """
+ Iterate only the leaf nodes in the tree.
+
+ >>> t = BinaryTree()
+ >>> t.insert(50)
+ >>> t.insert(75)
+ >>> t.insert(25)
+ >>> t.insert(66)
+ >>> t.insert(22)
+ >>> t.insert(13)
+
+ >>> for value in t.iterate_nodes_by_depth(2):
+ ... print(value)
+ 22
+ 66
+
+ >>> for value in t.iterate_nodes_by_depth(3):
+ ... print(value)
+ 13
+
+ """
+ if self.root is not None:
+ yield from self._iterate_by_depth(self.root, depth)
+
+ def _depth(self, node: Node, sofar: int) -> int:
+ depth_left = sofar + 1
+ depth_right = sofar + 1
+ if node.left is not None:
+ depth_left = self._depth(node.left, sofar + 1)
+ if node.right is not None:
+ depth_right = self._depth(node.right, sofar + 1)
+ return max(depth_left, depth_right)
+
+ def depth(self):
+ """
+ Returns the max height (depth) of the tree in plies (edge distance
+ from root).
+
+ >>> t = BinaryTree()
+ >>> t.depth()
+ 0
+
+ >>> t.insert(50)
+ >>> t.depth()
+ 1
+
+ >>> t.insert(65)
+ >>> t.depth()
+ 2
+
+ >>> t.insert(33)
+ >>> t.depth()
+ 2
+
+ >>> t.insert(2)
+ >>> t.insert(1)
+ >>> t.depth()
+ 4
+
+ """
+ if self.root is None:
+ return 0
+ return self._depth(self.root, 0)
+
+ def height(self):
+ return self.depth()
+
+
+if __name__ == '__main__':
+ import doctest
+ doctest.testmod()