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