github.com/wangyougui/gf/v2@v2.6.5/container/gtree/gtree_redblacktree.go (about) 1 // Copyright GoFrame Author(https://goframe.org). All Rights Reserved. 2 // 3 // This Source Code Form is subject to the terms of the MIT License. 4 // If a copy of the MIT was not distributed with this file, 5 // You can obtain one at https://github.com/wangyougui/gf. 6 7 package gtree 8 9 import ( 10 "bytes" 11 "fmt" 12 13 "github.com/wangyougui/gf/v2/container/gvar" 14 "github.com/wangyougui/gf/v2/internal/json" 15 "github.com/wangyougui/gf/v2/internal/rwmutex" 16 "github.com/wangyougui/gf/v2/util/gconv" 17 "github.com/wangyougui/gf/v2/util/gutil" 18 ) 19 20 type color bool 21 22 const ( 23 black, red color = true, false 24 ) 25 26 // RedBlackTree holds elements of the red-black tree. 27 type RedBlackTree struct { 28 mu rwmutex.RWMutex 29 root *RedBlackTreeNode 30 size int 31 comparator func(v1, v2 interface{}) int 32 } 33 34 // RedBlackTreeNode is a single element within the tree. 35 type RedBlackTreeNode struct { 36 Key interface{} 37 Value interface{} 38 color color 39 left *RedBlackTreeNode 40 right *RedBlackTreeNode 41 parent *RedBlackTreeNode 42 } 43 44 // NewRedBlackTree instantiates a red-black tree with the custom key comparator. 45 // The parameter `safe` is used to specify whether using tree in concurrent-safety, 46 // which is false in default. 47 func NewRedBlackTree(comparator func(v1, v2 interface{}) int, safe ...bool) *RedBlackTree { 48 return &RedBlackTree{ 49 mu: rwmutex.Create(safe...), 50 comparator: comparator, 51 } 52 } 53 54 // NewRedBlackTreeFrom instantiates a red-black tree with the custom key comparator and `data` map. 55 // The parameter `safe` is used to specify whether using tree in concurrent-safety, 56 // which is false in default. 57 func NewRedBlackTreeFrom(comparator func(v1, v2 interface{}) int, data map[interface{}]interface{}, safe ...bool) *RedBlackTree { 58 tree := NewRedBlackTree(comparator, safe...) 59 for k, v := range data { 60 tree.doSet(k, v) 61 } 62 return tree 63 } 64 65 // SetComparator sets/changes the comparator for sorting. 66 func (tree *RedBlackTree) SetComparator(comparator func(a, b interface{}) int) { 67 tree.mu.Lock() 68 defer tree.mu.Unlock() 69 tree.comparator = comparator 70 if tree.size > 0 { 71 data := make(map[interface{}]interface{}, tree.size) 72 tree.doIteratorAsc(tree.leftNode(), func(key, value interface{}) bool { 73 data[key] = value 74 return true 75 }) 76 // Resort the tree if comparator is changed. 77 tree.root = nil 78 tree.size = 0 79 for k, v := range data { 80 tree.doSet(k, v) 81 } 82 } 83 } 84 85 // Clone returns a new tree with a copy of current tree. 86 func (tree *RedBlackTree) Clone() *RedBlackTree { 87 newTree := NewRedBlackTree(tree.comparator, tree.mu.IsSafe()) 88 newTree.Sets(tree.Map()) 89 return newTree 90 } 91 92 // Set inserts key-value item into the tree. 93 func (tree *RedBlackTree) Set(key interface{}, value interface{}) { 94 tree.mu.Lock() 95 defer tree.mu.Unlock() 96 tree.doSet(key, value) 97 } 98 99 // Sets batch sets key-values to the tree. 100 func (tree *RedBlackTree) Sets(data map[interface{}]interface{}) { 101 tree.mu.Lock() 102 defer tree.mu.Unlock() 103 for k, v := range data { 104 tree.doSet(k, v) 105 } 106 } 107 108 // doSet inserts key-value item into the tree without mutex. 109 func (tree *RedBlackTree) doSet(key interface{}, value interface{}) { 110 insertedNode := (*RedBlackTreeNode)(nil) 111 if tree.root == nil { 112 // Assert key is of comparator's type for initial tree 113 tree.getComparator()(key, key) 114 tree.root = &RedBlackTreeNode{Key: key, Value: value, color: red} 115 insertedNode = tree.root 116 } else { 117 node := tree.root 118 loop := true 119 for loop { 120 compare := tree.getComparator()(key, node.Key) 121 switch { 122 case compare == 0: 123 // node.Key = key 124 node.Value = value 125 return 126 case compare < 0: 127 if node.left == nil { 128 node.left = &RedBlackTreeNode{Key: key, Value: value, color: red} 129 insertedNode = node.left 130 loop = false 131 } else { 132 node = node.left 133 } 134 case compare > 0: 135 if node.right == nil { 136 node.right = &RedBlackTreeNode{Key: key, Value: value, color: red} 137 insertedNode = node.right 138 loop = false 139 } else { 140 node = node.right 141 } 142 } 143 } 144 insertedNode.parent = node 145 } 146 tree.insertCase1(insertedNode) 147 tree.size++ 148 } 149 150 // Get searches the node in the tree by `key` and returns its value or nil if key is not found in tree. 151 func (tree *RedBlackTree) Get(key interface{}) (value interface{}) { 152 value, _ = tree.Search(key) 153 return 154 } 155 156 // doSetWithLockCheck checks whether value of the key exists with mutex.Lock, 157 // if not exists, set value to the map with given `key`, 158 // or else just return the existing value. 159 // 160 // When setting value, if `value` is type of <func() interface {}>, 161 // it will be executed with mutex.Lock of the hash map, 162 // and its return value will be set to the map with `key`. 163 // 164 // It returns value with given `key`. 165 func (tree *RedBlackTree) doSetWithLockCheck(key interface{}, value interface{}) interface{} { 166 tree.mu.Lock() 167 defer tree.mu.Unlock() 168 if node, found := tree.doSearch(key); found { 169 return node.Value 170 } 171 if f, ok := value.(func() interface{}); ok { 172 value = f() 173 } 174 if value != nil { 175 tree.doSet(key, value) 176 } 177 return value 178 } 179 180 // GetOrSet returns the value by key, 181 // or sets value with given `value` if it does not exist and then returns this value. 182 func (tree *RedBlackTree) GetOrSet(key interface{}, value interface{}) interface{} { 183 if v, ok := tree.Search(key); !ok { 184 return tree.doSetWithLockCheck(key, value) 185 } else { 186 return v 187 } 188 } 189 190 // GetOrSetFunc returns the value by key, 191 // or sets value with returned value of callback function `f` if it does not exist 192 // and then returns this value. 193 func (tree *RedBlackTree) GetOrSetFunc(key interface{}, f func() interface{}) interface{} { 194 if v, ok := tree.Search(key); !ok { 195 return tree.doSetWithLockCheck(key, f()) 196 } else { 197 return v 198 } 199 } 200 201 // GetOrSetFuncLock returns the value by key, 202 // or sets value with returned value of callback function `f` if it does not exist 203 // and then returns this value. 204 // 205 // GetOrSetFuncLock differs with GetOrSetFunc function is that it executes function `f` 206 // with mutex.Lock of the hash map. 207 func (tree *RedBlackTree) GetOrSetFuncLock(key interface{}, f func() interface{}) interface{} { 208 if v, ok := tree.Search(key); !ok { 209 return tree.doSetWithLockCheck(key, f) 210 } else { 211 return v 212 } 213 } 214 215 // GetVar returns a gvar.Var with the value by given `key`. 216 // The returned gvar.Var is un-concurrent safe. 217 func (tree *RedBlackTree) GetVar(key interface{}) *gvar.Var { 218 return gvar.New(tree.Get(key)) 219 } 220 221 // GetVarOrSet returns a gvar.Var with result from GetVarOrSet. 222 // The returned gvar.Var is un-concurrent safe. 223 func (tree *RedBlackTree) GetVarOrSet(key interface{}, value interface{}) *gvar.Var { 224 return gvar.New(tree.GetOrSet(key, value)) 225 } 226 227 // GetVarOrSetFunc returns a gvar.Var with result from GetOrSetFunc. 228 // The returned gvar.Var is un-concurrent safe. 229 func (tree *RedBlackTree) GetVarOrSetFunc(key interface{}, f func() interface{}) *gvar.Var { 230 return gvar.New(tree.GetOrSetFunc(key, f)) 231 } 232 233 // GetVarOrSetFuncLock returns a gvar.Var with result from GetOrSetFuncLock. 234 // The returned gvar.Var is un-concurrent safe. 235 func (tree *RedBlackTree) GetVarOrSetFuncLock(key interface{}, f func() interface{}) *gvar.Var { 236 return gvar.New(tree.GetOrSetFuncLock(key, f)) 237 } 238 239 // SetIfNotExist sets `value` to the map if the `key` does not exist, and then returns true. 240 // It returns false if `key` exists, and `value` would be ignored. 241 func (tree *RedBlackTree) SetIfNotExist(key interface{}, value interface{}) bool { 242 if !tree.Contains(key) { 243 tree.doSetWithLockCheck(key, value) 244 return true 245 } 246 return false 247 } 248 249 // SetIfNotExistFunc sets value with return value of callback function `f`, and then returns true. 250 // It returns false if `key` exists, and `value` would be ignored. 251 func (tree *RedBlackTree) SetIfNotExistFunc(key interface{}, f func() interface{}) bool { 252 if !tree.Contains(key) { 253 tree.doSetWithLockCheck(key, f()) 254 return true 255 } 256 return false 257 } 258 259 // SetIfNotExistFuncLock sets value with return value of callback function `f`, and then returns true. 260 // It returns false if `key` exists, and `value` would be ignored. 261 // 262 // SetIfNotExistFuncLock differs with SetIfNotExistFunc function is that 263 // it executes function `f` with mutex.Lock of the hash map. 264 func (tree *RedBlackTree) SetIfNotExistFuncLock(key interface{}, f func() interface{}) bool { 265 if !tree.Contains(key) { 266 tree.doSetWithLockCheck(key, f) 267 return true 268 } 269 return false 270 } 271 272 // Contains checks whether `key` exists in the tree. 273 func (tree *RedBlackTree) Contains(key interface{}) bool { 274 _, ok := tree.Search(key) 275 return ok 276 } 277 278 // doRemove removes the node from the tree by `key` without mutex. 279 func (tree *RedBlackTree) doRemove(key interface{}) (value interface{}) { 280 child := (*RedBlackTreeNode)(nil) 281 node, found := tree.doSearch(key) 282 if !found { 283 return 284 } 285 value = node.Value 286 if node.left != nil && node.right != nil { 287 p := node.left.maximumNode() 288 node.Key = p.Key 289 node.Value = p.Value 290 node = p 291 } 292 if node.left == nil || node.right == nil { 293 if node.right == nil { 294 child = node.left 295 } else { 296 child = node.right 297 } 298 if node.color == black { 299 node.color = tree.nodeColor(child) 300 tree.deleteCase1(node) 301 } 302 tree.replaceNode(node, child) 303 if node.parent == nil && child != nil { 304 child.color = black 305 } 306 } 307 tree.size-- 308 return 309 } 310 311 // Remove removes the node from the tree by `key`. 312 func (tree *RedBlackTree) Remove(key interface{}) (value interface{}) { 313 tree.mu.Lock() 314 defer tree.mu.Unlock() 315 return tree.doRemove(key) 316 } 317 318 // Removes batch deletes values of the tree by `keys`. 319 func (tree *RedBlackTree) Removes(keys []interface{}) { 320 tree.mu.Lock() 321 defer tree.mu.Unlock() 322 for _, key := range keys { 323 tree.doRemove(key) 324 } 325 } 326 327 // IsEmpty returns true if tree does not contain any nodes. 328 func (tree *RedBlackTree) IsEmpty() bool { 329 return tree.Size() == 0 330 } 331 332 // Size returns number of nodes in the tree. 333 func (tree *RedBlackTree) Size() int { 334 tree.mu.RLock() 335 defer tree.mu.RUnlock() 336 return tree.size 337 } 338 339 // Keys returns all keys in asc order. 340 func (tree *RedBlackTree) Keys() []interface{} { 341 var ( 342 keys = make([]interface{}, tree.Size()) 343 index = 0 344 ) 345 tree.IteratorAsc(func(key, value interface{}) bool { 346 keys[index] = key 347 index++ 348 return true 349 }) 350 return keys 351 } 352 353 // Values returns all values in asc order based on the key. 354 func (tree *RedBlackTree) Values() []interface{} { 355 var ( 356 values = make([]interface{}, tree.Size()) 357 index = 0 358 ) 359 tree.IteratorAsc(func(key, value interface{}) bool { 360 values[index] = value 361 index++ 362 return true 363 }) 364 return values 365 } 366 367 // Map returns all key-value items as map. 368 func (tree *RedBlackTree) Map() map[interface{}]interface{} { 369 m := make(map[interface{}]interface{}, tree.Size()) 370 tree.IteratorAsc(func(key, value interface{}) bool { 371 m[key] = value 372 return true 373 }) 374 return m 375 } 376 377 // MapStrAny returns all key-value items as map[string]interface{}. 378 func (tree *RedBlackTree) MapStrAny() map[string]interface{} { 379 m := make(map[string]interface{}, tree.Size()) 380 tree.IteratorAsc(func(key, value interface{}) bool { 381 m[gconv.String(key)] = value 382 return true 383 }) 384 return m 385 } 386 387 // Left returns the left-most (min) node or nil if tree is empty. 388 func (tree *RedBlackTree) Left() *RedBlackTreeNode { 389 tree.mu.RLock() 390 defer tree.mu.RUnlock() 391 node := tree.leftNode() 392 if tree.mu.IsSafe() { 393 return &RedBlackTreeNode{ 394 Key: node.Key, 395 Value: node.Value, 396 } 397 } 398 return node 399 } 400 401 // Right returns the right-most (max) node or nil if tree is empty. 402 func (tree *RedBlackTree) Right() *RedBlackTreeNode { 403 tree.mu.RLock() 404 defer tree.mu.RUnlock() 405 node := tree.rightNode() 406 if tree.mu.IsSafe() { 407 return &RedBlackTreeNode{ 408 Key: node.Key, 409 Value: node.Value, 410 } 411 } 412 return node 413 } 414 415 // leftNode returns the left-most (min) node or nil if tree is empty. 416 func (tree *RedBlackTree) leftNode() *RedBlackTreeNode { 417 p := (*RedBlackTreeNode)(nil) 418 n := tree.root 419 for n != nil { 420 p = n 421 n = n.left 422 } 423 return p 424 } 425 426 // rightNode returns the right-most (max) node or nil if tree is empty. 427 func (tree *RedBlackTree) rightNode() *RedBlackTreeNode { 428 p := (*RedBlackTreeNode)(nil) 429 n := tree.root 430 for n != nil { 431 p = n 432 n = n.right 433 } 434 return p 435 } 436 437 // Floor Finds floor node of the input key, return the floor node or nil if no floor node is found. 438 // Second return parameter is true if floor was found, otherwise false. 439 // 440 // Floor node is defined as the largest node that its key is smaller than or equal to the given `key`. 441 // A floor node may not be found, either because the tree is empty, or because 442 // all nodes in the tree are larger than the given node. 443 func (tree *RedBlackTree) Floor(key interface{}) (floor *RedBlackTreeNode, found bool) { 444 tree.mu.RLock() 445 defer tree.mu.RUnlock() 446 n := tree.root 447 for n != nil { 448 compare := tree.getComparator()(key, n.Key) 449 switch { 450 case compare == 0: 451 return n, true 452 case compare < 0: 453 n = n.left 454 case compare > 0: 455 floor, found = n, true 456 n = n.right 457 } 458 } 459 if found { 460 return 461 } 462 return nil, false 463 } 464 465 // Ceiling finds ceiling node of the input key, return the ceiling node or nil if no ceiling node is found. 466 // Second return parameter is true if ceiling was found, otherwise false. 467 // 468 // Ceiling node is defined as the smallest node that its key is larger than or equal to the given `key`. 469 // A ceiling node may not be found, either because the tree is empty, or because 470 // all nodes in the tree are smaller than the given node. 471 func (tree *RedBlackTree) Ceiling(key interface{}) (ceiling *RedBlackTreeNode, found bool) { 472 tree.mu.RLock() 473 defer tree.mu.RUnlock() 474 n := tree.root 475 for n != nil { 476 compare := tree.getComparator()(key, n.Key) 477 switch { 478 case compare == 0: 479 return n, true 480 case compare > 0: 481 n = n.right 482 case compare < 0: 483 ceiling, found = n, true 484 n = n.left 485 } 486 } 487 if found { 488 return 489 } 490 return nil, false 491 } 492 493 // Iterator is alias of IteratorAsc. 494 func (tree *RedBlackTree) Iterator(f func(key, value interface{}) bool) { 495 tree.IteratorAsc(f) 496 } 497 498 // IteratorFrom is alias of IteratorAscFrom. 499 func (tree *RedBlackTree) IteratorFrom(key interface{}, match bool, f func(key, value interface{}) bool) { 500 tree.IteratorAscFrom(key, match, f) 501 } 502 503 // IteratorAsc iterates the tree readonly in ascending order with given callback function `f`. 504 // If `f` returns true, then it continues iterating; or false to stop. 505 func (tree *RedBlackTree) IteratorAsc(f func(key, value interface{}) bool) { 506 tree.mu.RLock() 507 defer tree.mu.RUnlock() 508 tree.doIteratorAsc(tree.leftNode(), f) 509 } 510 511 // IteratorAscFrom iterates the tree readonly in ascending order with given callback function `f`. 512 // The parameter `key` specifies the start entry for iterating. The `match` specifies whether 513 // starting iterating if the `key` is fully matched, or else using index searching iterating. 514 // If `f` returns true, then it continues iterating; or false to stop. 515 func (tree *RedBlackTree) IteratorAscFrom(key interface{}, match bool, f func(key, value interface{}) bool) { 516 tree.mu.RLock() 517 defer tree.mu.RUnlock() 518 node, found := tree.doSearch(key) 519 if match { 520 if found { 521 tree.doIteratorAsc(node, f) 522 } 523 } else { 524 tree.doIteratorAsc(node, f) 525 } 526 } 527 528 func (tree *RedBlackTree) doIteratorAsc(node *RedBlackTreeNode, f func(key, value interface{}) bool) { 529 loop: 530 if node == nil { 531 return 532 } 533 if !f(node.Key, node.Value) { 534 return 535 } 536 if node.right != nil { 537 node = node.right 538 for node.left != nil { 539 node = node.left 540 } 541 goto loop 542 } 543 if node.parent != nil { 544 old := node 545 for node.parent != nil { 546 node = node.parent 547 if tree.getComparator()(old.Key, node.Key) <= 0 { 548 goto loop 549 } 550 } 551 } 552 } 553 554 // IteratorDesc iterates the tree readonly in descending order with given callback function `f`. 555 // If `f` returns true, then it continues iterating; or false to stop. 556 func (tree *RedBlackTree) IteratorDesc(f func(key, value interface{}) bool) { 557 tree.mu.RLock() 558 defer tree.mu.RUnlock() 559 tree.doIteratorDesc(tree.rightNode(), f) 560 } 561 562 // IteratorDescFrom iterates the tree readonly in descending order with given callback function `f`. 563 // The parameter `key` specifies the start entry for iterating. The `match` specifies whether 564 // starting iterating if the `key` is fully matched, or else using index searching iterating. 565 // If `f` returns true, then it continues iterating; or false to stop. 566 func (tree *RedBlackTree) IteratorDescFrom(key interface{}, match bool, f func(key, value interface{}) bool) { 567 tree.mu.RLock() 568 defer tree.mu.RUnlock() 569 node, found := tree.doSearch(key) 570 if match { 571 if found { 572 tree.doIteratorDesc(node, f) 573 } 574 } else { 575 tree.doIteratorDesc(node, f) 576 } 577 } 578 579 func (tree *RedBlackTree) doIteratorDesc(node *RedBlackTreeNode, f func(key, value interface{}) bool) { 580 loop: 581 if node == nil { 582 return 583 } 584 if !f(node.Key, node.Value) { 585 return 586 } 587 if node.left != nil { 588 node = node.left 589 for node.right != nil { 590 node = node.right 591 } 592 goto loop 593 } 594 if node.parent != nil { 595 old := node 596 for node.parent != nil { 597 node = node.parent 598 if tree.getComparator()(old.Key, node.Key) >= 0 { 599 goto loop 600 } 601 } 602 } 603 } 604 605 // Clear removes all nodes from the tree. 606 func (tree *RedBlackTree) Clear() { 607 tree.mu.Lock() 608 defer tree.mu.Unlock() 609 tree.root = nil 610 tree.size = 0 611 } 612 613 // Replace the data of the tree with given `data`. 614 func (tree *RedBlackTree) Replace(data map[interface{}]interface{}) { 615 tree.mu.Lock() 616 defer tree.mu.Unlock() 617 tree.root = nil 618 tree.size = 0 619 for k, v := range data { 620 tree.doSet(k, v) 621 } 622 } 623 624 // String returns a string representation of container. 625 func (tree *RedBlackTree) String() string { 626 if tree == nil { 627 return "" 628 } 629 tree.mu.RLock() 630 defer tree.mu.RUnlock() 631 str := "" 632 if tree.size != 0 { 633 tree.output(tree.root, "", true, &str) 634 } 635 return str 636 } 637 638 // Print prints the tree to stdout. 639 func (tree *RedBlackTree) Print() { 640 fmt.Println(tree.String()) 641 } 642 643 // Search searches the tree with given `key`. 644 // Second return parameter `found` is true if key was found, otherwise false. 645 func (tree *RedBlackTree) Search(key interface{}) (value interface{}, found bool) { 646 tree.mu.RLock() 647 defer tree.mu.RUnlock() 648 node, found := tree.doSearch(key) 649 if found { 650 return node.Value, true 651 } 652 return nil, false 653 } 654 655 // Flip exchanges key-value of the tree to value-key. 656 // Note that you should guarantee the value is the same type as key, 657 // or else the comparator would panic. 658 // 659 // If the type of value is different with key, you pass the new `comparator`. 660 func (tree *RedBlackTree) Flip(comparator ...func(v1, v2 interface{}) int) { 661 t := (*RedBlackTree)(nil) 662 if len(comparator) > 0 { 663 t = NewRedBlackTree(comparator[0], tree.mu.IsSafe()) 664 } else { 665 t = NewRedBlackTree(tree.comparator, tree.mu.IsSafe()) 666 } 667 tree.IteratorAsc(func(key, value interface{}) bool { 668 t.doSet(value, key) 669 return true 670 }) 671 tree.mu.Lock() 672 tree.root = t.root 673 tree.size = t.size 674 tree.mu.Unlock() 675 } 676 677 func (tree *RedBlackTree) output(node *RedBlackTreeNode, prefix string, isTail bool, str *string) { 678 if node.right != nil { 679 newPrefix := prefix 680 if isTail { 681 newPrefix += "│ " 682 } else { 683 newPrefix += " " 684 } 685 tree.output(node.right, newPrefix, false, str) 686 } 687 *str += prefix 688 if isTail { 689 *str += "└── " 690 } else { 691 *str += "┌── " 692 } 693 *str += fmt.Sprintf("%v\n", node.Key) 694 if node.left != nil { 695 newPrefix := prefix 696 if isTail { 697 newPrefix += " " 698 } else { 699 newPrefix += "│ " 700 } 701 tree.output(node.left, newPrefix, true, str) 702 } 703 } 704 705 // doSearch searches the tree with given `key` without mutex. 706 // It returns the node if found or otherwise nil. 707 func (tree *RedBlackTree) doSearch(key interface{}) (node *RedBlackTreeNode, found bool) { 708 node = tree.root 709 for node != nil { 710 compare := tree.getComparator()(key, node.Key) 711 switch { 712 case compare == 0: 713 return node, true 714 case compare < 0: 715 node = node.left 716 case compare > 0: 717 node = node.right 718 } 719 } 720 return node, false 721 } 722 723 func (node *RedBlackTreeNode) grandparent() *RedBlackTreeNode { 724 if node != nil && node.parent != nil { 725 return node.parent.parent 726 } 727 return nil 728 } 729 730 func (node *RedBlackTreeNode) uncle() *RedBlackTreeNode { 731 if node == nil || node.parent == nil || node.parent.parent == nil { 732 return nil 733 } 734 return node.parent.sibling() 735 } 736 737 func (node *RedBlackTreeNode) sibling() *RedBlackTreeNode { 738 if node == nil || node.parent == nil { 739 return nil 740 } 741 if node == node.parent.left { 742 return node.parent.right 743 } 744 return node.parent.left 745 } 746 747 func (tree *RedBlackTree) rotateLeft(node *RedBlackTreeNode) { 748 right := node.right 749 tree.replaceNode(node, right) 750 node.right = right.left 751 if right.left != nil { 752 right.left.parent = node 753 } 754 right.left = node 755 node.parent = right 756 } 757 758 func (tree *RedBlackTree) rotateRight(node *RedBlackTreeNode) { 759 left := node.left 760 tree.replaceNode(node, left) 761 node.left = left.right 762 if left.right != nil { 763 left.right.parent = node 764 } 765 left.right = node 766 node.parent = left 767 } 768 769 func (tree *RedBlackTree) replaceNode(old *RedBlackTreeNode, new *RedBlackTreeNode) { 770 if old.parent == nil { 771 tree.root = new 772 } else { 773 if old == old.parent.left { 774 old.parent.left = new 775 } else { 776 old.parent.right = new 777 } 778 } 779 if new != nil { 780 new.parent = old.parent 781 } 782 } 783 784 func (tree *RedBlackTree) insertCase1(node *RedBlackTreeNode) { 785 if node.parent == nil { 786 node.color = black 787 } else { 788 tree.insertCase2(node) 789 } 790 } 791 792 func (tree *RedBlackTree) insertCase2(node *RedBlackTreeNode) { 793 if tree.nodeColor(node.parent) == black { 794 return 795 } 796 tree.insertCase3(node) 797 } 798 799 func (tree *RedBlackTree) insertCase3(node *RedBlackTreeNode) { 800 uncle := node.uncle() 801 if tree.nodeColor(uncle) == red { 802 node.parent.color = black 803 uncle.color = black 804 node.grandparent().color = red 805 tree.insertCase1(node.grandparent()) 806 } else { 807 tree.insertCase4(node) 808 } 809 } 810 811 func (tree *RedBlackTree) insertCase4(node *RedBlackTreeNode) { 812 grandparent := node.grandparent() 813 if node == node.parent.right && node.parent == grandparent.left { 814 tree.rotateLeft(node.parent) 815 node = node.left 816 } else if node == node.parent.left && node.parent == grandparent.right { 817 tree.rotateRight(node.parent) 818 node = node.right 819 } 820 tree.insertCase5(node) 821 } 822 823 func (tree *RedBlackTree) insertCase5(node *RedBlackTreeNode) { 824 node.parent.color = black 825 grandparent := node.grandparent() 826 grandparent.color = red 827 if node == node.parent.left && node.parent == grandparent.left { 828 tree.rotateRight(grandparent) 829 } else if node == node.parent.right && node.parent == grandparent.right { 830 tree.rotateLeft(grandparent) 831 } 832 } 833 834 func (node *RedBlackTreeNode) maximumNode() *RedBlackTreeNode { 835 if node == nil { 836 return nil 837 } 838 for node.right != nil { 839 return node.right 840 } 841 return node 842 } 843 844 func (tree *RedBlackTree) deleteCase1(node *RedBlackTreeNode) { 845 if node.parent == nil { 846 return 847 } 848 tree.deleteCase2(node) 849 } 850 851 func (tree *RedBlackTree) deleteCase2(node *RedBlackTreeNode) { 852 sibling := node.sibling() 853 if tree.nodeColor(sibling) == red { 854 node.parent.color = red 855 sibling.color = black 856 if node == node.parent.left { 857 tree.rotateLeft(node.parent) 858 } else { 859 tree.rotateRight(node.parent) 860 } 861 } 862 tree.deleteCase3(node) 863 } 864 865 func (tree *RedBlackTree) deleteCase3(node *RedBlackTreeNode) { 866 sibling := node.sibling() 867 if tree.nodeColor(node.parent) == black && 868 tree.nodeColor(sibling) == black && 869 tree.nodeColor(sibling.left) == black && 870 tree.nodeColor(sibling.right) == black { 871 sibling.color = red 872 tree.deleteCase1(node.parent) 873 } else { 874 tree.deleteCase4(node) 875 } 876 } 877 878 func (tree *RedBlackTree) deleteCase4(node *RedBlackTreeNode) { 879 sibling := node.sibling() 880 if tree.nodeColor(node.parent) == red && 881 tree.nodeColor(sibling) == black && 882 tree.nodeColor(sibling.left) == black && 883 tree.nodeColor(sibling.right) == black { 884 sibling.color = red 885 node.parent.color = black 886 } else { 887 tree.deleteCase5(node) 888 } 889 } 890 891 func (tree *RedBlackTree) deleteCase5(node *RedBlackTreeNode) { 892 sibling := node.sibling() 893 if node == node.parent.left && 894 tree.nodeColor(sibling) == black && 895 tree.nodeColor(sibling.left) == red && 896 tree.nodeColor(sibling.right) == black { 897 sibling.color = red 898 sibling.left.color = black 899 tree.rotateRight(sibling) 900 } else if node == node.parent.right && 901 tree.nodeColor(sibling) == black && 902 tree.nodeColor(sibling.right) == red && 903 tree.nodeColor(sibling.left) == black { 904 sibling.color = red 905 sibling.right.color = black 906 tree.rotateLeft(sibling) 907 } 908 tree.deleteCase6(node) 909 } 910 911 func (tree *RedBlackTree) deleteCase6(node *RedBlackTreeNode) { 912 sibling := node.sibling() 913 sibling.color = tree.nodeColor(node.parent) 914 node.parent.color = black 915 if node == node.parent.left && tree.nodeColor(sibling.right) == red { 916 sibling.right.color = black 917 tree.rotateLeft(node.parent) 918 } else if tree.nodeColor(sibling.left) == red { 919 sibling.left.color = black 920 tree.rotateRight(node.parent) 921 } 922 } 923 924 func (tree *RedBlackTree) nodeColor(node *RedBlackTreeNode) color { 925 if node == nil { 926 return black 927 } 928 return node.color 929 } 930 931 // MarshalJSON implements the interface MarshalJSON for json.Marshal. 932 func (tree RedBlackTree) MarshalJSON() (jsonBytes []byte, err error) { 933 if tree.root == nil { 934 return []byte("null"), nil 935 } 936 buffer := bytes.NewBuffer(nil) 937 buffer.WriteByte('{') 938 tree.Iterator(func(key, value interface{}) bool { 939 valueBytes, valueJsonErr := json.Marshal(value) 940 if valueJsonErr != nil { 941 err = valueJsonErr 942 return false 943 } 944 if buffer.Len() > 1 { 945 buffer.WriteByte(',') 946 } 947 buffer.WriteString(fmt.Sprintf(`"%v":%s`, key, valueBytes)) 948 return true 949 }) 950 buffer.WriteByte('}') 951 return buffer.Bytes(), nil 952 } 953 954 // UnmarshalJSON implements the interface UnmarshalJSON for json.Unmarshal. 955 func (tree *RedBlackTree) UnmarshalJSON(b []byte) error { 956 tree.mu.Lock() 957 defer tree.mu.Unlock() 958 if tree.comparator == nil { 959 tree.comparator = gutil.ComparatorString 960 } 961 var data map[string]interface{} 962 if err := json.UnmarshalUseNumber(b, &data); err != nil { 963 return err 964 } 965 for k, v := range data { 966 tree.doSet(k, v) 967 } 968 return nil 969 } 970 971 // UnmarshalValue is an interface implement which sets any type of value for map. 972 func (tree *RedBlackTree) UnmarshalValue(value interface{}) (err error) { 973 tree.mu.Lock() 974 defer tree.mu.Unlock() 975 if tree.comparator == nil { 976 tree.comparator = gutil.ComparatorString 977 } 978 for k, v := range gconv.Map(value) { 979 tree.doSet(k, v) 980 } 981 return 982 } 983 984 // getComparator returns the comparator if it's previously set, 985 // or else it panics. 986 func (tree *RedBlackTree) getComparator() func(a, b interface{}) int { 987 if tree.comparator == nil { 988 panic("comparator is missing for tree") 989 } 990 return tree.comparator 991 }