| 42 | self.rightChild.parent = self |
| 43 | |
| 44 | class BinarySearchTree: |
| 45 | def __init__(self): |
| 46 | self.root = None |
| 47 | self.size = 0 |
| 48 | |
| 49 | def length(self): |
| 50 | return self.size |
| 51 | |
| 52 | def __len__(self): |
| 53 | return self.size |
| 54 | |
| 55 | def __iter__(self): |
| 56 | return self.root.__iter__() |
| 57 | |
| 58 | def put(self, key, val): |
| 59 | if self.root: |
| 60 | self._put(key, val, self.root) |
| 61 | else: |
| 62 | self.root = TreeNode(key, val) |
| 63 | self.size = self.size + 1 |
| 64 | |
| 65 | def _put(self, key, val, currentNode): |
| 66 | if key < currentNode.key: |
| 67 | if currentNode.hasLeftChild(): |
| 68 | self._put(key, val, currentNode.leftChild) |
| 69 | else: |
| 70 | currentNode.leftChild = TreeNode(key, val, parent=currentNode) |
| 71 | self.updateBalance(currentNode.leftChild) |
| 72 | else: |
| 73 | if currentNode.hasRightChild(): |
| 74 | self._put(key, val, currentNode.rightChild) |
| 75 | else: |
| 76 | currentNode.rightChild = TreeNode(key, val, parent=currentNode) |
| 77 | self.updateBalance(currentNode.rightChild) |
| 78 | |
| 79 | def updateBalance(self, node): |
| 80 | if node.balanceFactor > 1 or node.balanceFactor < -1: |
| 81 | self.rebalance(node) |
| 82 | return |
| 83 | if node.parent != None: |
| 84 | if node.isLeftChild(): |
| 85 | node.parent.balanceFactor += 1 |
| 86 | elif node.isRightChild(): |
| 87 | node.parent.balanceFactor -= 1 |
| 88 | |
| 89 | if node.parent.balanceFactor != 0: |
| 90 | self.updateBalance(node.parent) |
| 91 | |
| 92 | def rotateLeft(self, rotRoot): |
| 93 | newRoot = rotRoot.rightChild |
| 94 | rotRoot.rightChild = newRoot.leftChild |
| 95 | if newRoot.leftChild != None: |
| 96 | newRoot.leftChild.parent = rotRoot |
| 97 | newRoot.parent = rotRoot.parent |
| 98 | if rotRoot.isRoot(): |
| 99 | self.root = newRoot |
| 100 | else: |
| 101 | if rotRoot.isLeftChild(): |