github.com/djordje200179/extendedlibrary/datastructures@v1.7.1-0.20240227175559-d09520a92dd4/maps/rbt/balancing.go (about)

     1  package rbt
     2  
     3  func (tree *Tree[K, V]) fixInsert(node *Node[K, V]) {
     4  	if node.parent == nil {
     5  		node.color = black
     6  
     7  		return
     8  	} else if node.parent.color == black {
     9  		return
    10  	}
    11  
    12  	uncle := node.parent.Sibling()
    13  	if nodeColor(uncle) == red {
    14  		node.parent.color = black
    15  		uncle.color = black
    16  		node.parent.parent.color = red
    17  		tree.fixInsert(node.parent.parent)
    18  
    19  		return
    20  	}
    21  
    22  	grandparent := node.parent.parent
    23  	if node == node.parent.rightChild && node.parent == grandparent.leftChild {
    24  		tree.rotateLeft(node.parent)
    25  		node = node.leftChild
    26  	} else if node == node.parent.leftChild && node.parent == grandparent.rightChild {
    27  		tree.rotateRight(node.parent)
    28  		node = node.rightChild
    29  	}
    30  
    31  	node.parent.color = black
    32  	grandparent.color = red
    33  
    34  	if node == node.parent.leftChild && node.parent == grandparent.leftChild {
    35  		tree.rotateRight(grandparent)
    36  	} else {
    37  		tree.rotateLeft(grandparent)
    38  	}
    39  }
    40  
    41  func (tree *Tree[K, V]) fixRemove(node *Node[K, V]) {
    42  	for node != tree.root && nodeColor(node) == black {
    43  		if node == node.parent.leftChild {
    44  			sibling := node.parent.rightChild
    45  
    46  			if nodeColor(sibling) == red {
    47  				sibling.color = black
    48  				node.parent.color = red
    49  				tree.rotateLeft(node.parent)
    50  				sibling = node.parent.rightChild
    51  			}
    52  
    53  			if nodeColor(sibling.leftChild) == black && nodeColor(sibling.rightChild) == black {
    54  				sibling.color = red
    55  				node = node.parent
    56  			} else {
    57  				if nodeColor(sibling.rightChild) == black {
    58  					sibling.leftChild.color = black
    59  					sibling.color = red
    60  					tree.rotateRight(sibling)
    61  					sibling = node.parent.rightChild
    62  				}
    63  
    64  				sibling.color = node.parent.color
    65  				node.parent.color = black
    66  				sibling.rightChild.color = black
    67  				tree.rotateLeft(node.parent)
    68  				node = tree.root
    69  			}
    70  		} else {
    71  			sibling := node.parent.leftChild
    72  
    73  			if nodeColor(sibling) == red {
    74  				sibling.color = black
    75  				node.parent.color = red
    76  				tree.rotateRight(node.parent)
    77  				sibling = node.parent.leftChild
    78  			}
    79  
    80  			if nodeColor(sibling.leftChild) == black && nodeColor(sibling.rightChild) == black {
    81  				sibling.color = red
    82  				node = node.parent
    83  			} else {
    84  				if nodeColor(sibling.leftChild) == black {
    85  					sibling.rightChild.color = black
    86  					sibling.color = red
    87  					tree.rotateLeft(sibling)
    88  					sibling = node.parent.leftChild
    89  				}
    90  
    91  				sibling.color = node.parent.color
    92  				node.parent.color = black
    93  				sibling.leftChild.color = black
    94  				tree.rotateRight(node.parent)
    95  				node = tree.root
    96  			}
    97  		}
    98  	}
    99  
   100  	node.color = black
   101  }
   102  
   103  func (tree *Tree[K, V]) rotateLeft(node *Node[K, V]) {
   104  	rightNode := node.rightChild
   105  	node.rightChild = rightNode.leftChild
   106  	if rightNode.leftChild != nil {
   107  		rightNode.leftChild.parent = node
   108  	}
   109  
   110  	rightNode.parent = node.parent
   111  	if node.parent == nil {
   112  		tree.root = rightNode
   113  	} else if node == node.parent.leftChild {
   114  		node.parent.leftChild = rightNode
   115  	} else {
   116  		node.parent.rightChild = rightNode
   117  	}
   118  
   119  	rightNode.leftChild = node
   120  	node.parent = rightNode
   121  }
   122  
   123  func (tree *Tree[K, V]) rotateRight(node *Node[K, V]) {
   124  	leftNode := node.leftChild
   125  	node.leftChild = leftNode.rightChild
   126  	if leftNode.rightChild != nil {
   127  		leftNode.rightChild.parent = node
   128  	}
   129  
   130  	leftNode.parent = node.parent
   131  	if node.parent == nil {
   132  		tree.root = leftNode
   133  	} else if node == node.parent.rightChild {
   134  		node.parent.rightChild = leftNode
   135  	} else {
   136  		node.parent.leftChild = leftNode
   137  	}
   138  	leftNode.rightChild = node
   139  	node.parent = leftNode
   140  }