AVL 树是最早被发明的自平衡二叉查找树。在 AVL 树中，任意节点对应的两颗子树的最大高度差为 1，因此他也被称为高度平衡树。查找、插入和删除在平均和最坏的情况下的时间复杂度都是 O(log n)。增加和删除元素的操作则节能需要一次或多次树旋转，以实现树的重新平衡。

结构定义

type AVLTreeNode struct {
	Left   *AVLTreeNode
	Right  *AVLTreeNode
	Val    int
	Height int // 当前节点深度 (左子树与右子树的差)
}

旋转树节点

// 将树进行左旋
func ll(node *AVLTreeNode) *AVLTreeNode {
	nodeR := node.Right
	node.Right = nodeR.Left
	nodeR.Left = node
	node.Height = max(height(node.Left), height(node.Right)) + 1
	nodeR.Height = max(height(nodeR.Right), node.Height) + 1
	return nodeR
}

// 将树进行右旋
func rr(node *AVLTreeNode) *AVLTreeNode {
	nodeL := node.Left
	node.Left = nodeL.Right
	nodeL.Right = node
	node.Height = max(height(node.Left), height(node.Right)) + 1
	nodeL.Height = max(height(nodeL.Left), node.Height) + 1
	return nodeL
}

 // 将树进行右旋后左旋
func lr(node *AVLTreeNode) *AVLTreeNode {
	node.Left = ll(node.Left)
	return rr(node)
}

// 将树进行左旋后右旋
func rl(node *AVLTreeNode) *AVLTreeNode {
	node.Right = rr(node.Right)
	return ll(node)
}

 // 若树失去平衡性，`进行旋转
func handleBF(node *AVLTreeNode) *AVLTreeNode {
        // 左子树失衡
	if height(node.Left)-height(node.Right) == 2 {
		if height(node.Left.Left)-height(node.Left.Right) > 0 {
			node = rr(node)  // 左子树的左子树失衡
		} else {
			node = lr(node) // 左子树的右子树失衡
		}
            // 右子树失衡
	} else if height(node.Right)-height(node.Left) == 2 {
		if height(node.Right.Left)-height(node.Right.Right) > 0 {
			node = rl(node)  // 右子树的左子树失衡
		} else {
			node = ll(node)  // 右子树的右子树失衡
		}
	}
        // 重新计算当前节点深度
	node.Height = max(height(node.Left), height(node.Right)) + 1
	return node
}

添加元素

func height(t *AVLTreeNode) int {
	if t == nil {
		return -1
	}
	return t.Height
}

func max(i, j int) int {
	if i < j {
		return j
	}
	return i
}

func insert(node *AVLTreeNode, val int) *AVLTreeNode {
	if node == nil {
		node = NewAVLTree(val)
	} else if val < node.Val {
		node.Left = insert(node.Left, val)
		node = handleBF(node)
	} else if val > node.Val {
		node.Right = insert(node.Right, val)
		node = handleBF(node)
	}

	return node
}

删除节点

func delete(node *AVLTreeNode, val int) *AVLTreeNode {
	if node == nil {
		fmt.Println("没有元素")
		return nil
	} else if node.Val < val {
		node.Right = delete(node.Right, val)
		node = handleBF(node)
	} else if val < node.Val {
		node.Left = delete(node.Left, val)
		node = handleBF(node)
	} else if val == node.Val {
		// 若待删除节点为叶子节点, 直接删除
		if node.Left == nil && node.Right == nil {
			node = nil
			return node
		// 若待删除节点有左子树或右子树，左子树或右子树代替被删除节点
		} else if node.Left != nil && node.Right == nil {
			nodeL := node.Left
			node.Left = nil
			return handleBF(nodeL)
		} else if node.Right != nil && node.Left == nil {
			nodeR := node.Right
			node.Right = nil
			return handleBF(nodeR)
		// 若待删除节点有左子树和右子树
		} else {
			// 找出待删除节点的后继节点和后继节点的父节点
			p, n := searchNode(node.Right)
			// 若后继节点是待删除节点的子节点，后继节点代替待删除节点
			if node.Right == n {
				n.Left = node.Left
				node = nil
				return handleBF(n)
		        // 若后继节点有右子树，后继节点的父节点的左子树为后继节点的右子树
			} else if n.Right != nil {
				p.Left = n.Right
				n.Left = node.Left
				n.Right = node.Right
				node = nil
				return handleBF(n)
		        // 若后继节点没有右子树，后继节点代替待删除节点
			} else if n.Right == nil {
				n.Left = node.Left
				n.Right = node.Right
				node = nil
				p.Left = nil
				return handleBF(n)
			}
		}
	}
	return node
}

找出后继节点与后继节点的父节点

func searchNode(node *AVLTreeNode) (*AVLTreeNode, *AVLTreeNode) {
	if node.Left == nil {
		return nil, node
	}
	for {
		if node.Left.Left != nil {
			node = node.Left
		} else {
			return node, node.Left
		}
	}
}