collect/bst.py

   1 #!/usr/bin/env python3
   2
   3 from typing import Any, Optional
   4
   5
   6 class Node(object):
   7     def __init__(self, value: Any) -> None:
   8         self.left = None
   9         self.right = None
  10         self.value = value
  11
  12
  13 class BinaryTree(object):
  14     def __init__(self):
  15         self.root = None
  16         self.count = 0
  17         self.traverse = None
  18
  19     def get_root(self) -> Optional[Node]:
  20         return self.root
  21
  22     def insert(self, value: Any):
  23         """
  24         Insert something into the tree.
  25
  26         >>> t = BinaryTree()
  27         >>> t.insert(10)
  28         >>> t.insert(20)
  29         >>> t.insert(5)
  30         >>> len(t)
  31         3
  32
  33         >>> t.get_root().value
  34         10
  35
  36         """
  37         if self.root is None:
  38             self.root = Node(value)
  39             self.count = 1
  40         else:
  41             self._insert(value, self.root)
  42
  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)
  48             else:
  49                 node.left = Node(value)
  50                 self.count += 1
  51         else:
  52             if node.right is not None:
  53                 self._insert(value, node.right)
  54             else:
  55                 node.right = Node(value)
  56                 self.count += 1
  57
  58     def __getitem__(self, value: Any) -> Optional[Node]:
  59         """
  60         Find an item in the tree and return its Node.  Returns
  61         None if the item is not in the tree.
  62
  63         >>> t = BinaryTree()
  64         >>> t[99]
  65
  66         >>> t.insert(10)
  67         >>> t.insert(20)
  68         >>> t.insert(5)
  69         >>> t[10].value
  70         10
  71
  72         >>> t[99]
  73
  74         """
  75         if self.root is not None:
  76             return self._find(value, self.root)
  77         return None
  78
  79     def _find(self, value: Any, node: Node) -> Optional[Node]:
  80         """Find helper"""
  81         if value == node.value:
  82             return node
  83         elif (value < node.value and node.left is not None):
  84             return self._find(value, node.left)
  85         else:
  86             assert value > node.value
  87             if node.right is not None:
  88                 return self._find(value, node.right)
  89         return None
  90
  91     def __delitem__(self, value: Any) -> bool:
  92         """
  93         Delete an item from the tree and preserve the BST property.
  94
  95                             50
  96                            /  \
  97                          25    75
  98                         /     /  \
  99                       22    66    85
 100                      /
 101                    13
 102
 103
 104         >>> t = BinaryTree()
 105         >>> t.insert(50)
 106         >>> t.insert(75)
 107         >>> t.insert(25)
 108         >>> t.insert(66)
 109         >>> t.insert(22)
 110         >>> t.insert(13)
 111         >>> t.insert(85)
 112
 113         >>> for value in t.iterate_inorder():
 114         ...     print(value)
 115         13
 116         22
 117         25
 118         50
 119         66
 120         75
 121         85
 122
 123         >>> del t[22]  # Note: bool result is discarded
 124
 125         >>> for value in t.iterate_inorder():
 126         ...     print(value)
 127         13
 128         25
 129         50
 130         66
 131         75
 132         85
 133
 134         >>> t.__delitem__(13)
 135         True
 136         >>> for value in t.iterate_inorder():
 137         ...     print(value)
 138         25
 139         50
 140         66
 141         75
 142         85
 143
 144         >>> t.__delitem__(75)
 145         True
 146         >>> for value in t.iterate_inorder():
 147         ...     print(value)
 148         25
 149         50
 150         66
 151         85
 152
 153         >>> t.__delitem__(99)
 154         False
 155
 156         """
 157         if self.root is not None:
 158             ret = self._delete(value, None, self.root)
 159             if ret:
 160                 self.count -= 1
 161                 if self.count == 0:
 162                     self.root = None
 163             return ret
 164         return False
 165
 166     def _delete(self, value: Any, parent: Optional[Node], node: Node) -> bool:
 167         """Delete helper"""
 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:
 173                         parent.left = None
 174                     else:
 175                         assert parent.right == node
 176                         parent.right = None
 177                 return True
 178
 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
 185                     else:
 186                         assert parent.right == node
 187                         parent.right = node.right
 188                 return True
 189
 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
 196                     else:
 197                         assert parent.right == node
 198                         parent.right = node.left
 199                 return True
 200
 201             # Node has both a left and right.
 202             else:
 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)
 213         return False
 214
 215     def __len__(self):
 216         """
 217         Returns the count of items in the tree.
 218
 219         >>> t = BinaryTree()
 220         >>> len(t)
 221         0
 222         >>> t.insert(50)
 223         >>> len(t)
 224         1
 225         >>> t.__delitem__(50)
 226         True
 227         >>> len(t)
 228         0
 229         >>> t.insert(75)
 230         >>> t.insert(25)
 231         >>> t.insert(66)
 232         >>> t.insert(22)
 233         >>> t.insert(13)
 234         >>> t.insert(85)
 235         >>> len(t)
 236         6
 237
 238         """
 239         return self.count
 240
 241     def __contains__(self, value: Any) -> bool:
 242         """
 243         Returns True if the item is in the tree; False otherwise.
 244
 245         """
 246         return self.__getitem__(value) is not None
 247
 248     def _iterate_preorder(self, node: Node):
 249         yield node.value
 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)
 254
 255     def _iterate_inorder(self, node: Node):
 256         if node.left is not None:
 257             yield from self._iterate_inorder(node.left)
 258         yield node.value
 259         if node.right is not None:
 260             yield from self._iterate_inorder(node.right)
 261
 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)
 267         yield node.value
 268
 269     def iterate_preorder(self):
 270         """
 271         Yield the tree's items in a preorder traversal sequence.
 272
 273         >>> t = BinaryTree()
 274         >>> t.insert(50)
 275         >>> t.insert(75)
 276         >>> t.insert(25)
 277         >>> t.insert(66)
 278         >>> t.insert(22)
 279         >>> t.insert(13)
 280
 281         >>> for value in t.iterate_preorder():
 282         ...     print(value)
 283         50
 284         25
 285         22
 286         13
 287         75
 288         66
 289
 290         """
 291         if self.root is not None:
 292             yield from self._iterate_preorder(self.root)
 293
 294     def iterate_inorder(self):
 295         """
 296         Yield the tree's items in a preorder traversal sequence.
 297
 298         >>> t = BinaryTree()
 299         >>> t.insert(50)
 300         >>> t.insert(75)
 301         >>> t.insert(25)
 302         >>> t.insert(66)
 303         >>> t.insert(22)
 304         >>> t.insert(13)
 305
 306         >>> for value in t.iterate_inorder():
 307         ...     print(value)
 308         13
 309         22
 310         25
 311         50
 312         66
 313         75
 314
 315         """
 316         if self.root is not None:
 317             yield from self._iterate_inorder(self.root)
 318
 319     def iterate_postorder(self):
 320         """
 321         Yield the tree's items in a preorder traversal sequence.
 322
 323         >>> t = BinaryTree()
 324         >>> t.insert(50)
 325         >>> t.insert(75)
 326         >>> t.insert(25)
 327         >>> t.insert(66)
 328         >>> t.insert(22)
 329         >>> t.insert(13)
 330
 331         >>> for value in t.iterate_postorder():
 332         ...     print(value)
 333         13
 334         22
 335         25
 336         66
 337         75
 338         50
 339
 340         """
 341         if self.root is not None:
 342             yield from self._iterate_postorder(self.root)
 343
 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:
 350             yield node.value
 351
 352     def iterate_leaves(self):
 353         """
 354         Iterate only the leaf nodes in the tree.
 355
 356         >>> t = BinaryTree()
 357         >>> t.insert(50)
 358         >>> t.insert(75)
 359         >>> t.insert(25)
 360         >>> t.insert(66)
 361         >>> t.insert(22)
 362         >>> t.insert(13)
 363
 364         >>> for value in t.iterate_leaves():
 365         ...     print(value)
 366         13
 367         66
 368
 369         """
 370         if self.root is not None:
 371             yield from self._iterate_leaves(self.root)
 372
 373     def _iterate_by_depth(self, node: Node, depth: int):
 374         if depth == 0:
 375             yield node.value
 376         else:
 377             assert depth > 0
 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)
 382
 383     def iterate_nodes_by_depth(self, depth: int):
 384         """
 385         Iterate only the leaf nodes in the tree.
 386
 387         >>> t = BinaryTree()
 388         >>> t.insert(50)
 389         >>> t.insert(75)
 390         >>> t.insert(25)
 391         >>> t.insert(66)
 392         >>> t.insert(22)
 393         >>> t.insert(13)
 394
 395         >>> for value in t.iterate_nodes_by_depth(2):
 396         ...     print(value)
 397         22
 398         66
 399
 400         >>> for value in t.iterate_nodes_by_depth(3):
 401         ...     print(value)
 402         13
 403
 404         """
 405         if self.root is not None:
 406             yield from self._iterate_by_depth(self.root, depth)
 407
 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)
 416
 417     def depth(self):
 418         """
 419         Returns the max height (depth) of the tree in plies (edge distance
 420         from root).
 421
 422         >>> t = BinaryTree()
 423         >>> t.depth()
 424         0
 425
 426         >>> t.insert(50)
 427         >>> t.depth()
 428         1
 429
 430         >>> t.insert(65)
 431         >>> t.depth()
 432         2
 433
 434         >>> t.insert(33)
 435         >>> t.depth()
 436         2
 437
 438         >>> t.insert(2)
 439         >>> t.insert(1)
 440         >>> t.depth()
 441         4
 442
 443         """
 444         if self.root is None:
 445             return 0
 446         return self._depth(self.root, 0)
 447
 448     def height(self):
 449         return self.depth()
 450
 451
 452 if __name__ == '__main__':
 453     import doctest
 454     doctest.testmod()