github.com/benz9527/toy-box/algo@v0.0.0-20240221120937-66c0c6bd5abd/tree/avl_tree.go (about)

     1  package tree
     2  
     3  func CreateAVLNode(item int) *Node {
     4  	return &Node{Item: item, Height: 1, Left: nil, Right: nil}
     5  }
     6  
     7  func AVLHeight(n *Node) int {
     8  	if n == nil {
     9  		return 0
    10  	}
    11  	return n.Height
    12  }
    13  
    14  func Max(a, b int) int {
    15  	if a > b {
    16  		return a
    17  	}
    18  	return b
    19  }
    20  
    21  func LeftRotate(xn *Node) *Node {
    22  	yn := xn.Right
    23  	temp := yn.Left
    24  
    25  	xn.Right = temp
    26  	yn.Left = xn
    27  
    28  	xn.Height = Max(AVLHeight(xn.Left), AVLHeight(xn.Right)) + 1
    29  	yn.Height = Max(AVLHeight(yn.Left), AVLHeight(yn.Right)) + 1
    30  	return yn
    31  }
    32  
    33  func RightRotate(yn *Node) *Node {
    34  	xn := yn.Left
    35  	temp := xn.Right
    36  
    37  	yn.Left = temp
    38  	xn.Right = yn
    39  
    40  	xn.Height = Max(AVLHeight(xn.Left), AVLHeight(xn.Right)) + 1
    41  	yn.Height = Max(AVLHeight(yn.Left), AVLHeight(yn.Right)) + 1
    42  	return xn
    43  }
    44  
    45  func GetBalanceFactor(n *Node) int {
    46  	if n == nil {
    47  		return 0
    48  	}
    49  	return AVLHeight(n.Left) - AVLHeight(n.Right)
    50  }
    51  
    52  func AVLInsertNode(n *Node, newItem int) *Node {
    53  	if n == nil {
    54  		return CreateAVLNode(newItem)
    55  	}
    56  
    57  	if newItem < n.Item {
    58  		n.Left = AVLInsertNode(n.Left, newItem)
    59  	} else if newItem > n.Item {
    60  		n.Right = AVLInsertNode(n.Right, newItem)
    61  	} else {
    62  		return n
    63  	}
    64  
    65  	n.Height = 1 + Max(AVLHeight(n.Left), AVLHeight(n.Right))
    66  	bf := GetBalanceFactor(n)
    67  	if bf > 1 {
    68  		if newItem < n.Left.Item {
    69  			return RightRotate(n)
    70  		} else if newItem > n.Left.Item {
    71  			n.Left = LeftRotate(n.Left)
    72  			return RightRotate(n)
    73  		}
    74  	}
    75  
    76  	if bf < -1 {
    77  		if newItem > n.Right.Item {
    78  			return LeftRotate(n)
    79  		} else if newItem < n.Right.Item {
    80  			n.Right = RightRotate(n.Right)
    81  			return LeftRotate(n)
    82  		}
    83  	}
    84  	return n
    85  }
    86  
    87  func AVLMinimumValueNode(n *Node) *Node {
    88  	cur := n
    89  	for cur.Left != nil {
    90  		cur = cur.Left
    91  	}
    92  	return cur
    93  }
    94  
    95  func AVLDeleteNode(n *Node, deletedItem int) *Node {
    96  	if n == nil {
    97  		return n
    98  	}
    99  
   100  	if deletedItem < n.Item {
   101  		n.Left = AVLDeleteNode(n.Left, deletedItem)
   102  	} else if deletedItem > n.Item {
   103  		n.Right = AVLDeleteNode(n.Right, deletedItem)
   104  	} else {
   105  		// n will be removed, n is a leaf.
   106  		if n.Left == nil ||
   107  			n.Right == nil {
   108  			var temp *Node = nil
   109  
   110  			if temp == n.Left {
   111  				temp = n.Right
   112  			} else {
   113  				temp = n.Left
   114  			}
   115  
   116  			if temp == nil {
   117  				temp = n
   118  				n = nil
   119  			} else {
   120  				n = temp
   121  			}
   122  		} else {
   123  			temp := AVLMinimumValueNode(n.Right)
   124  			n.Item = temp.Item
   125  			n.Right = AVLDeleteNode(n.Right, temp.Item)
   126  		}
   127  	}
   128  
   129  	if n == nil {
   130  		return n
   131  	}
   132  
   133  	n.Height = Max(AVLHeight(n.Left), AVLHeight(n.Right)) + 1
   134  	bf := GetBalanceFactor(n)
   135  	if bf > 1 {
   136  		if GetBalanceFactor(n.Left) >= 0 {
   137  			return RightRotate(n)
   138  		} else {
   139  			n.Left = LeftRotate(n.Left)
   140  			return RightRotate(n)
   141  		}
   142  	}
   143  
   144  	if bf < -1 {
   145  		if GetBalanceFactor(n.Right) <= 0 {
   146  			return LeftRotate(n)
   147  		} else {
   148  			n.Right = RightRotate(n.Right)
   149  			return LeftRotate(n)
   150  		}
   151  	}
   152  	return n
   153  }