Branch node with single array optimization.
| 122 | |
| 123 | |
| 124 | class OptimizedBranchNode: |
| 125 | """Branch node with single array optimization.""" |
| 126 | |
| 127 | def __init__(self, capacity: int): |
| 128 | self.capacity = capacity |
| 129 | self.num_keys = 0 |
| 130 | # Array layout: keys[0:capacity], children[capacity:capacity*2+1] |
| 131 | self.data = [None] * (capacity * 2 + 1) |
| 132 | |
| 133 | def is_leaf(self) -> bool: |
| 134 | return False |
| 135 | |
| 136 | def find_child_index(self, key) -> int: |
| 137 | """Binary search for child index.""" |
| 138 | return bisect.bisect_right(self.data, key, 0, self.num_keys) |
| 139 | |
| 140 | def get_child(self, key): |
| 141 | """Get child node for given key.""" |
| 142 | index = self.find_child_index(key) |
| 143 | return self.data[self.capacity + index] |
| 144 | |
| 145 | def set_child(self, index: int, child): |
| 146 | """Set child at index.""" |
| 147 | self.data[self.capacity + index] = child |
| 148 | |
| 149 | def insert(self, key, right_child) -> Optional[Tuple[Any, "OptimizedBranchNode"]]: |
| 150 | """Insert key and right child.""" |
| 151 | pos = bisect.bisect_left(self.data, key, 0, self.num_keys) |
| 152 | |
| 153 | # Check if split needed |
| 154 | if self.num_keys >= self.capacity: |
| 155 | return self._split_and_insert(pos, key, right_child) |
| 156 | |
| 157 | # Shift keys and children |
| 158 | if pos < self.num_keys: |
| 159 | # Shift keys |
| 160 | self.data[pos + 1 : self.num_keys + 1] = self.data[pos : self.num_keys] |
| 161 | # Shift children (one extra child) |
| 162 | start_child = self.capacity + pos + 1 |
| 163 | end_child = self.capacity + self.num_keys + 1 |
| 164 | self.data[start_child + 1 : end_child + 1] = self.data[ |
| 165 | start_child:end_child |
| 166 | ] |
| 167 | |
| 168 | # Insert |
| 169 | self.data[pos] = key |
| 170 | self.data[self.capacity + pos + 1] = right_child |
| 171 | self.num_keys += 1 |
| 172 | return None |
| 173 | |
| 174 | def _split_and_insert( |
| 175 | self, pos: int, key, right_child |
| 176 | ) -> Tuple[Any, "OptimizedBranchNode"]: |
| 177 | """Split branch node.""" |
| 178 | new_node = OptimizedBranchNode(self.capacity) |
| 179 | mid = self.capacity // 2 |
| 180 | |
| 181 | # Collect all keys and children |
no outgoing calls