stl_treemap.tree
Red-black tree implementation.
1"""Red-black tree implementation.""" 2 3from __future__ import annotations 4 5from collections.abc import Callable, Collection 6from typing import Any 7 8from stl_treemap.insertion_result import InsertionResult 9from stl_treemap.iterators import ReverseIterator, TreeIterator 10from stl_treemap.policies import KeyOnlyPolicy, PolicyInterface, ValueOnlyPolicy 11from stl_treemap.py_iterators import PyIterator, PyReverseIterator 12from stl_treemap.tree_node import Head, NodeColors, TreeNode 13 14_INSERT_MULTI = 1 15_INSERT_UNIQUE = 2 16_INSERT_REPLACE = 3 17 18 19class Tree[K, V](Collection[K]): 20 """@private Red-black tree.""" 21 22 def __init__(self) -> None: 23 """Construct an empty tree.""" 24 self.head: Head[K, V] = Head() 25 self.compare: Callable[[Any, Any], int] = Tree.compare 26 self.value_policy: PolicyInterface[K, V] = KeyOnlyPolicy() 27 28 @staticmethod 29 def compare(lhs: Any, rhs: Any) -> int: 30 """ 31 3-way comparison, similar to strcmp in C. 32 33 Args: 34 lhs: Left-hand side value. 35 rhs: Right-hand side value. 36 37 Returns: 38 -1 if lhs < rhs, 0 if equal, 1 if lhs > rhs. 39 40 """ 41 if lhs < rhs: 42 return -1 43 if lhs == rhs: 44 return 0 45 return 1 46 47 def clear(self) -> None: 48 """Delete all nodes in the tree.""" 49 self.head = Head() 50 51 def size(self) -> int: 52 """@private Return the number of nodes in the tree.""" 53 return self.head.size 54 55 def compare_nodes(self, lhs: TreeNode[K, V], rhs: TreeNode[K, V]) -> int: 56 """@private Call 3-way comparison on node keys.""" 57 return self.compare(lhs.key, rhs.key) 58 59 def replace_node(self, old_node: TreeNode[K, V], new_node: TreeNode[K, V] | None) -> None: 60 """@private Update parent and child pointers when swapping nodes during rotations.""" 61 if old_node is new_node: 62 return 63 if old_node.parent is None: 64 self.head.root = new_node 65 elif old_node is old_node.parent.left: 66 old_node.parent.left = new_node 67 else: 68 old_node.parent.right = new_node 69 if not self.is_leaf(new_node): 70 new_node.parent = old_node.parent 71 72 def rotate_left(self, node: TreeNode[K, V]) -> None: 73 r""" 74 Rebalance the tree with a left rotation. 75 76 :: 77 78 X Y 79 / \ rotate / \ 80 Y c -------> a X 81 / \ / \ 82 a b b c 83 84 Args: 85 node: Root of the sub-tree to rotate. 86 87 Raises: 88 ValueError: When the node has no right child (tree is corrupted). 89 90 """ 91 right = node.right 92 if self.is_leaf(right): 93 raise ValueError("rotateLeft can't be performed. The tree is corrupted") 94 self.replace_node(node, right) 95 node.right = right.left 96 if right.left is not None: 97 right.left.parent = node 98 right.left = node 99 node.parent = right 100 101 def rotate_right(self, node: TreeNode[K, V]) -> None: 102 """ 103 Rebalance the tree with a right rotation (mirror of rotate_left). 104 105 Args: 106 node: Root of the sub-tree to rotate. 107 108 Raises: 109 ValueError: When the node has no left child (tree is corrupted). 110 111 """ 112 left = node.left 113 if self.is_leaf(left): 114 raise ValueError("rotateRight can't be performed. The tree is corrupted") 115 self.replace_node(node, left) 116 node.left = left.right 117 if left.right is not None: 118 left.right.parent = node 119 left.right = node 120 node.parent = left 121 122 def is_leaf(self, node: Any) -> bool: 123 """ 124 Return True for nodes that should not be traversed. 125 126 Args: 127 node: Node to inspect. 128 129 Returns: 130 True for None and the head sentinel; False for all real nodes. 131 132 """ 133 return node is None or node is self.head 134 135 def fetch_color(self, node: Any) -> int: 136 """ 137 Return the color of a node, treating leaves as BLACK. 138 139 Args: 140 node: Node to inspect. 141 142 Returns: 143 NodeColors.BLACK for leaf nodes; the node's stored color otherwise. 144 145 """ 146 if self.is_leaf(node): 147 return NodeColors.BLACK 148 return node.color 149 150 def is_black(self, node: Any) -> bool: 151 """Return True when the node color is BLACK.""" 152 return self.fetch_color(node) == NodeColors.BLACK 153 154 def is_red(self, node: Any) -> bool: 155 """Return True when the node color is RED.""" 156 return self.fetch_color(node) == NodeColors.RED 157 158 # === INSERT === 159 160 def insert_multi(self, node: TreeNode[K, V]) -> InsertionResult[TreeIterator[K, V]]: 161 """ 162 Insert a node, allowing duplicate keys. 163 164 Args: 165 node: Node to insert. 166 167 Returns: 168 InsertionResult indicating whether the node was added and an iterator to it. 169 170 """ 171 return self.insert_node(node, _INSERT_MULTI) 172 173 def insert_unique(self, node: TreeNode[K, V]) -> InsertionResult[TreeIterator[K, V]]: 174 """ 175 Insert a node only if no node with the same key exists. 176 177 Args: 178 node: Node to insert. 179 180 Returns: 181 InsertionResult indicating whether the node was added and an iterator to it. 182 183 """ 184 return self.insert_node(node, _INSERT_UNIQUE) 185 186 def insert_or_replace(self, node: TreeNode[K, V]) -> InsertionResult[TreeIterator[K, V]]: 187 """ 188 Insert a node, replacing the value of any existing node with the same key. 189 190 Args: 191 node: Node to insert. 192 193 Returns: 194 InsertionResult indicating whether the node was added or replaced and an iterator to it. 195 196 """ 197 return self.insert_node(node, _INSERT_REPLACE) 198 199 def insert_node(self, n: TreeNode[K, V], mode: int = _INSERT_MULTI) -> InsertionResult[TreeIterator[K, V]]: 200 """ 201 @private Insert a node, update the head sentinel, and rebalance the tree. 202 203 Args: 204 n: Node to insert. 205 mode: One of _INSERT_MULTI, _INSERT_UNIQUE, or _INSERT_REPLACE. 206 207 Returns: 208 InsertionResult indicating outcome and providing an iterator to the node. 209 210 """ 211 res = self.insert_node_internal(self.head.root, n, mode) 212 if res.was_added: 213 if self.head.size == 0: 214 self.head.root = n 215 self.head.leftmost = n 216 self.head.rightmost = n 217 n.left = self.head 218 n.right = self.head 219 elif self.head.leftmost.left is n: 220 self.head.leftmost = n 221 n.left = self.head 222 elif self.head.rightmost.right is n: 223 self.head.rightmost = n 224 n.right = self.head 225 self.insert_repair_tree(n) 226 self.head.size += 1 227 return res 228 229 def insert_node_internal(self, root: Any, n: TreeNode[K, V], mode: int) -> InsertionResult[TreeIterator[K, V]]: 230 """ 231 @private Descend the tree and place the node according to mode. 232 233 Args: 234 root: Root node of the tree. 235 n: Node to insert. 236 mode: One of _INSERT_MULTI, _INSERT_UNIQUE, or _INSERT_REPLACE. 237 238 Returns: 239 InsertionResult indicating outcome and providing an iterator to the node. 240 241 """ 242 x = root 243 y = None 244 rc = -1 245 while not self.is_leaf(x): 246 y = x 247 rc = self.compare_nodes(n, y) 248 if rc < 0: 249 x = y.left 250 elif rc > 0: 251 x = y.right 252 else: 253 if mode == _INSERT_UNIQUE: 254 return InsertionResult(False, False, None) 255 if mode == _INSERT_REPLACE: 256 self.value_policy.copy(y, n) 257 return InsertionResult(False, True, TreeIterator(y, self)) 258 x = y.right # _INSERT_MULTI: continue right to allow duplicates 259 if self.is_leaf(y): 260 n.parent = None 261 n.left = self.head 262 n.right = self.head 263 else: 264 n.parent = y 265 if rc < 0: 266 y.left = n 267 else: 268 y.right = n 269 return InsertionResult(True, False, TreeIterator(n, self)) 270 271 def insert_repair_tree(self, n: TreeNode[K, V]) -> None: 272 """ 273 @private Restore red-black invariants after insertion. 274 275 See https://en.wikipedia.org/wiki/Red-black_tree#Insertion 276 277 Args: 278 n: Newly inserted node. 279 280 """ 281 if n.parent is None: 282 self.repair_case1(n) 283 elif self.is_black(n.parent): 284 pass # case 2: parent is black — no fix needed 285 elif self.is_red(n.uncle()): 286 self.repair_case3(n) 287 else: 288 self.repair_case4(n) 289 290 def repair_case1(self, n: TreeNode[K, V]) -> None: 291 """@private https://en.wikipedia.org/wiki/Red-black_tree#Insertion.""" 292 n.color = NodeColors.BLACK 293 294 def repair_case3(self, n: TreeNode[K, V]) -> None: 295 """@private https://en.wikipedia.org/wiki/Red-black_tree#Insertion.""" 296 n.parent.color = NodeColors.BLACK 297 n.uncle().color = NodeColors.BLACK 298 n.grandparent().color = NodeColors.RED 299 self.insert_repair_tree(n.grandparent()) 300 301 def repair_case4(self, node: TreeNode[K, V]) -> None: 302 """@private https://en.wikipedia.org/wiki/Red-black_tree#Insertion.""" 303 n = node 304 p = n.parent 305 g = n.grandparent() 306 if not self.is_leaf(g.left) and n is g.left.right: 307 self.rotate_left(p) 308 n = n.left 309 elif not self.is_leaf(g.right) and n is g.right.left: 310 self.rotate_right(p) 311 n = n.right 312 p = n.parent 313 g = n.grandparent() 314 if n is p.left: 315 self.rotate_right(g) 316 else: 317 self.rotate_left(g) 318 p.color = NodeColors.BLACK 319 g.color = NodeColors.RED 320 321 # === FETCH MAX / MIN === 322 323 def fetch_maximum(self, node: TreeNode[K, V]) -> TreeNode[K, V]: 324 """ 325 Return the node with the largest key in the subtree. 326 327 Args: 328 node: Root of the subtree to search. 329 330 Returns: 331 Node with the highest key. 332 333 """ 334 res = node 335 while not self.is_leaf(res.right): 336 res = res.right 337 return res 338 339 def fetch_minimum(self, node: TreeNode[K, V]) -> TreeNode[K, V]: 340 """ 341 Return the node with the smallest key in the subtree. 342 343 Args: 344 node: Root of the subtree to search. 345 346 Returns: 347 Node with the lowest key. 348 349 """ 350 res = node 351 while not self.is_leaf(res.left): 352 res = res.left 353 return res 354 355 # === ERASE === 356 357 def erase(self, node: TreeNode[K, V]) -> None: 358 """ 359 Remove a node from the tree. 360 361 Args: 362 node: Node to remove; ignored if it is a leaf. 363 364 """ 365 if self.is_leaf(node): 366 return 367 self.erase_internal(node) 368 self.head.size -= 1 369 370 def erase_internal(self, node_param: TreeNode[K, V]) -> None: 371 """ 372 @private https://en.wikipedia.org/wiki/Red-black_tree#Removal. 373 374 Args: 375 node_param: Node to remove. 376 377 """ 378 node = node_param 379 if not self.is_leaf(node.left) and not self.is_leaf(node.right): 380 pred = self.fetch_maximum(node.left) 381 self.value_policy.copy(node, pred) 382 node = pred 383 384 child = node.left if self.is_leaf(node.right) else node.right 385 386 if self.is_black(node): 387 self.erase_case1(node) 388 self.replace_node(node, child) 389 if not self.is_leaf(child) and child.parent is None: 390 # child becomes root 391 child.color = NodeColors.BLACK 392 393 h = self.head 394 if self.is_leaf(child): 395 if h.leftmost is node: 396 p = node.parent 397 if p is None: 398 h.leftmost = h 399 else: 400 h.leftmost = p 401 p.left = h 402 if h.rightmost is node: 403 p = node.parent 404 if p is None: 405 h.rightmost = h 406 else: 407 h.rightmost = p 408 p.right = h 409 else: 410 if h.leftmost is node: 411 h.leftmost = child 412 child.left = h 413 if h.rightmost is node: 414 h.rightmost = child 415 child.right = h 416 417 def erase_case1(self, node: TreeNode[K, V]) -> None: 418 """@private https://en.wikipedia.org/wiki/Red-black_tree#Removal.""" 419 if node.parent is not None: 420 self.erase_case2(node) 421 422 def erase_case2(self, node: TreeNode[K, V]) -> None: 423 """@private https://en.wikipedia.org/wiki/Red-black_tree#Removal.""" 424 s = node.sibling() 425 if self.is_red(s): 426 node.parent.color = NodeColors.RED 427 s.color = NodeColors.BLACK 428 if node is node.parent.left: 429 self.rotate_left(node.parent) 430 else: 431 self.rotate_right(node.parent) 432 self.erase_case3(node) 433 434 def erase_case3(self, node: TreeNode[K, V]) -> None: 435 """@private https://en.wikipedia.org/wiki/Red-black_tree#Removal.""" 436 s = node.sibling() 437 p = node.parent 438 if self.is_black(p) and self.is_black(s) and self.is_black(s.left) and self.is_black(s.right): 439 s.color = NodeColors.RED 440 self.erase_case1(p) 441 else: 442 self.erase_case4(node) 443 444 def erase_case4(self, node: TreeNode[K, V]) -> None: 445 """@private https://en.wikipedia.org/wiki/Red-black_tree#Removal.""" 446 s = node.sibling() 447 p = node.parent 448 if self.is_red(p) and self.is_black(s) and self.is_black(s.left) and self.is_black(s.right): 449 s.color = NodeColors.RED 450 p.color = NodeColors.BLACK 451 else: 452 self.erase_case5(node) 453 454 def erase_case5(self, node: TreeNode[K, V]) -> None: 455 """@private https://en.wikipedia.org/wiki/Red-black_tree#Removal.""" 456 s = node.sibling() 457 p = node.parent 458 if node is p.left and self.is_red(s.left) and self.is_black(s.right): 459 s.color = NodeColors.RED 460 s.left.color = NodeColors.BLACK 461 self.rotate_right(s) 462 elif node is p.right and self.is_black(s.left) and self.is_red(s.right): 463 s.color = NodeColors.RED 464 s.right.color = NodeColors.BLACK 465 self.rotate_left(s) 466 self.erase_case6(node) 467 468 def erase_case6(self, node: TreeNode[K, V]) -> None: 469 """@private https://en.wikipedia.org/wiki/Red-black_tree#Removal.""" 470 s = node.sibling() 471 p = node.parent 472 s.color = self.fetch_color(p) 473 p.color = NodeColors.BLACK 474 if node is p.left: 475 s.right.color = NodeColors.BLACK 476 self.rotate_left(p) 477 else: 478 s.left.color = NodeColors.BLACK 479 self.rotate_right(p) 480 481 # === SEARCH === 482 483 def find(self, k: K) -> TreeIterator[K, V]: 484 """ 485 Return an iterator for the node with the given key. 486 487 Args: 488 k: Key to search for. 489 490 Returns: 491 Iterator pointing to the matching node, or end() if not found. 492 493 """ 494 x = self.head.root 495 while not self.is_leaf(x): 496 rc = self.compare(x.key, k) 497 if rc > 0: 498 x = x.left 499 elif rc < 0: 500 x = x.right 501 else: 502 return TreeIterator(x, self) 503 return TreeIterator(self.head, self) 504 505 def lower_bound(self, k: K) -> TreeIterator[K, V]: 506 """ 507 Return an iterator to the first node with key >= k. 508 509 Args: 510 k: Key to search for. 511 512 Returns: 513 Iterator to the first node not less than k, or end() if none exists. 514 515 """ 516 y: Any = self.head 517 x = y.root 518 while not self.is_leaf(x): 519 rc = self.compare(x.key, k) 520 if rc >= 0: 521 y = x 522 x = x.left 523 else: 524 x = x.right 525 return TreeIterator(y, self) 526 527 def upper_bound(self, k: K) -> TreeIterator[K, V]: 528 """ 529 Return an iterator to the first node with key > k. 530 531 Args: 532 k: Key to search for. 533 534 Returns: 535 Iterator to the first node greater than k, or end() if none exists. 536 537 """ 538 y: Any = self.head 539 x = y.root 540 while not self.is_leaf(x): 541 rc = self.compare(x.key, k) 542 if rc > 0: 543 y = x 544 x = x.left 545 else: 546 x = x.right 547 return TreeIterator(y, self) 548 549 # === STL ITERATORS === 550 551 def begin(self) -> TreeIterator[K, V]: 552 """Return an iterator to the node with the lowest key.""" 553 return TreeIterator(self.head.leftmost, self) 554 555 def end(self) -> TreeIterator[K, V]: 556 """Return an iterator to the position past the last node.""" 557 return TreeIterator(self.head, self) 558 559 def rbegin(self) -> ReverseIterator[K, V]: 560 """Return a reverse iterator to the node with the highest key.""" 561 return ReverseIterator(self.head.rightmost, self) 562 563 def rend(self) -> ReverseIterator[K, V]: 564 """Return a reverse iterator to the position before the first node.""" 565 return ReverseIterator(self.head, self) 566 567 # === PYTHON ITERATORS === 568 569 def py_begin(self) -> TreeNode[K, V]: 570 """@private Provide support for regular Python forward iteration.""" 571 return self.head.leftmost 572 573 def py_end(self) -> TreeNode[K, V]: 574 """@private Provide support for regular Python forward iteration.""" 575 return self.head 576 577 def py_rbegin(self) -> TreeNode[K, V]: 578 """@private Provide support for Python backward iteration.""" 579 return self.head.rightmost 580 581 def py_rend(self) -> TreeNode[K, V]: 582 """@private Provide support for Python backward iteration.""" 583 return self.head 584 585 def next(self, node: Any) -> Any: 586 """ 587 Return the next node in ascending key order. 588 589 Args: 590 node: Node to advance from. 591 592 Returns: 593 Successor node, or the head sentinel when past the last node. 594 595 """ 596 n = node 597 if n is self.head: 598 return self.head.leftmost 599 if n.right is self.head: 600 return self.head 601 if n.right is not None: 602 return self.fetch_minimum(n.right) 603 while n.parent.left is not n: 604 n = n.parent 605 return n.parent 606 607 def prev(self, node: Any) -> Any: 608 """ 609 Return the previous node in ascending key order. 610 611 Args: 612 node: Node to retreat from. 613 614 Returns: 615 Predecessor node, or the head sentinel when before the first node. 616 617 """ 618 n = node 619 if n is self.head: 620 return self.head.rightmost 621 if n.left is self.head: 622 return self.head 623 if n.left is not None: 624 return self.fetch_maximum(n.left) 625 while n.parent.right is not n: 626 n = n.parent 627 return n.parent 628 629 def items(self) -> PyIterator[tuple[K, V]]: 630 """ 631 Python forward iterator. 632 633 Returns: 634 Iterator for all elements in the tree in the order of keys 635 636 """ 637 return PyIterator(self, self.value_policy) 638 639 def backwards(self) -> PyReverseIterator[K]: 640 """ 641 Python reverse iterator. 642 643 Returns: 644 Iterator for all keys in the tree in the reverse order of keys 645 646 """ 647 return self.keys().backwards() 648 649 def first(self) -> K | V | tuple[K, V] | None: 650 """ 651 Return the first element, or None if the tree is empty. 652 653 Returns: 654 Value fetched by value_policy from the first node, or None. 655 656 """ 657 if self.size() == 0: 658 return None 659 return self.value_policy.fetch(self.begin().node) 660 661 def last(self) -> K | V | tuple[K, V] | None: 662 """ 663 Return the last element, or None if the tree is empty. 664 665 Returns: 666 Value fetched by value_policy from the last node, or None. 667 668 """ 669 if self.size() == 0: 670 return None 671 return self.value_policy.fetch(self.rbegin().node) 672 673 def __str__(self) -> str: 674 """Return a string representation of the container's contents.""" 675 parts = [] 676 it = self.begin() 677 end = self.end() 678 while not it.equals(end): 679 parts.append(self.value_policy.to_string(it.node)) 680 it.next() 681 return "{" + ",".join(parts) + "}" 682 683 def __repr__(self) -> str: 684 """Return a string representation of the container's contents.""" 685 return self.__str__() 686 687 def __iter__(self) -> PyIterator[K | V | tuple[K, V]]: 688 """Return forward key iterator.""" 689 return self.keys() 690 691 def __getitem__(self, key: K) -> V: 692 """ 693 Return the value associated with key. 694 695 Args: 696 key: Key to look up. 697 698 Returns: 699 Value stored at key. 700 701 Raises: 702 KeyError: When key is not present in the tree. 703 704 """ 705 it = self.find(key) 706 if self.is_leaf(it.node): 707 raise KeyError(key) 708 return it.node.value 709 710 def __setitem__(self, key: K, value: V) -> None: 711 """ 712 Insert or replace the value for key. 713 714 Args: 715 key: Key to insert or update. 716 value: Value to associate with the key. 717 718 """ 719 n: TreeNode[K, V] = TreeNode() 720 n.key = key 721 n.value = value 722 self.insert_or_replace(n) 723 724 def __contains__(self, k: K) -> bool: 725 """Return true if tree contains provided key.""" 726 iter = self.find(k) 727 return not iter.equals(self.end()) 728 729 def __delitem__(self, k: K): 730 """Delete item by key or raise KeyError.""" 731 iter = self.find(k) 732 if not iter.equals(self.end()): 733 self.erase(iter.node) 734 else: 735 raise KeyError(f"Key {k} not found") 736 737 def __len__(self) -> int: 738 """Return number of elements in the tree.""" 739 return self.head.size 740 741 def keys(self) -> PyIterator[K]: 742 """Yield keys for each node in ascending key order.""" 743 return PyIterator(self, KeyOnlyPolicy()) 744 745 def values(self) -> PyIterator[V]: 746 """Yield values for each node in ascending key order.""" 747 return PyIterator(self, ValueOnlyPolicy())