leetcode/avl/avlref.go

package avl

import "fmt"

// AVL 自平衡二叉搜索树 —— 正确实现参考
//
// 核心设计点（对比你原文件里的 4 个 bug）：
//  1. 节点维护 Height，平衡因子 BF = height(Left) - height(Right)；合法范围 {-1,0,1}
//  2. 旋转函数【返回新子树根】，调用方用返回值接回父节点的 Left/Right
//     —— 这解决了"旋转接不回去"的根本问题
//  3. 失衡判定看【子节点的 BF】，区分 LL/RR/LR/RL 四种 case，不是只看自己 BF
//  4. 不用全局 depthIncre flag，递归返回新根，height 直接从子节点算

type Node struct {
	Val    int
	Left   *Node
	Right  *Node
	Height int // 以该节点为根的子树高度；空节点=0，叶子=1
}

func height(n *Node) int {
	if n == nil {
		return 0
	}
	return n.Height
}

func balanceFactor(n *Node) int {
	if n == nil {
		return 0
	}
	return height(n.Left) - height(n.Right)
}

func updateHeight(n *Node) {
	l, r := height(n.Left), height(n.Right)
	if l > r {
		n.Height = l + 1
	} else {
		n.Height = r + 1
	}
}

// 左旋：右子太重时，把右子 y 提上来，x 沉为 y 的左子
//
//	  x                y
//	 / \              / \
//	A   y     →      x   C
//	   / \          / \
//	  B   C        A   B
//
// B 的值在 x、y 之间，必须留在中间 → 挂到 x 的右侧
func rotateLeft(x *Node) *Node {
	y := x.Right
	x.Right = y.Left
	y.Left = x
	updateHeight(x) // 先更新成为孩子的 x
	updateHeight(y) // 再更新成为根的 y
	return y
}

// 右旋：左子太重时，把左子 x 提上来，y 沉为 x 的右子
//
//	    y                x
//	   / \              / \
//	  x   C     →      A   y
//	 / \                  / \
//	A   B                B   C
//
// B 的值在 x、y 之间 → 挂到 y 的左侧
func rotateRight(y *Node) *Node {
	x := y.Left
	y.Left = x.Right
	x.Right = y
	updateHeight(y) // 先更新成为孩子的 y
	updateHeight(x) // 再更新成为根的 x
	return x
}

// 插入：递归返回新子树根
func Insert(root *Node, val int) *Node {
	if root == nil {
		return &Node{Val: val, Height: 1}
	}
	if val < root.Val {
		root.Left = Insert(root.Left, val)
	} else if val > root.Val {
		root.Right = Insert(root.Right, val)
	} else {
		return root // 不允许重复值
	}
	updateHeight(root)
	bf := balanceFactor(root)

	// LL: 左子左重（左子 BF >= 0）→ 右旋一次
	if bf > 1 && balanceFactor(root.Left) >= 0 {
		return rotateRight(root)
	}
	// LR: 左子右重（左子 BF < 0）→ 先左旋左子，再右旋自己
	if bf > 1 && balanceFactor(root.Left) < 0 {
		root.Left = rotateLeft(root.Left)
		return rotateRight(root)
	}
	// RR: 右子右重（右子 BF <= 0）→ 左旋一次
	if bf < -1 && balanceFactor(root.Right) <= 0 {
		return rotateLeft(root)
	}
	// RL: 右子左重（右子 BF > 0）→ 先右旋右子，再左旋自己
	if bf < -1 && balanceFactor(root.Right) > 0 {
		root.Right = rotateRight(root.Right)
		return rotateLeft(root)
	}
	return root
}

// 删除：递归返回新子树根
func Delete(root *Node, val int) *Node {
	if root == nil {
		return nil
	}
	if val < root.Val {
		root.Left = Delete(root.Left, val)
	} else if val > root.Val {
		root.Right = Delete(root.Right, val)
	} else {
		// 命中要删的节点
		if root.Left == nil {
			return root.Right
		}
		if root.Right == nil {
			return root.Left
		}
		// 两子都在：用右子树最小值替换当前值，再删右子树里的那个最小值
		minNode := root.Right
		for minNode.Left != nil {
			minNode = minNode.Left
		}
		root.Val = minNode.Val
		root.Right = Delete(root.Right, minNode.Val)
	}
	updateHeight(root)
	bf := balanceFactor(root)
	// 重平衡逻辑与 Insert 完全一致（看子节点 BF 判 case）
	if bf > 1 && balanceFactor(root.Left) >= 0 {
		return rotateRight(root)
	}
	if bf > 1 && balanceFactor(root.Left) < 0 {
		root.Left = rotateLeft(root.Left)
		return rotateRight(root)
	}
	if bf < -1 && balanceFactor(root.Right) <= 0 {
		return rotateLeft(root)
	}
	if bf < -1 && balanceFactor(root.Right) > 0 {
		root.Right = rotateRight(root.Right)
		return rotateLeft(root)
	}
	return root
}

func Search(root *Node, val int) bool {
	for root != nil {
		if val < root.Val {
			root = root.Left
		} else if val > root.Val {
			root = root.Right
		} else {
			return true
		}
	}
	return false
}

// 中序遍历 → 应得到升序序列（验证 BST 性质没被破坏）
func Inorder(root *Node, out *[]int) {
	if root == nil {
		return
	}
	Inorder(root.Left, out)
	*out = append(*out, root.Val)
	Inorder(root.Right, out)
}

func runAVLDemo() {
	var root *Node

	// 测试 1: 顺序插入 1..7（普通 BST 会退化成高度=7 的链表）
	for _, v := range []int{1, 2, 3, 4, 5, 6, 7} {
		root = Insert(root, v)
	}
	var arr []int
	Inorder(root, &arr)
	fmt.Println("顺序插入 1..7 中序遍历:", arr)
	fmt.Printf("树高 = %d (普通BST会=7, AVL 应 ≤ 4)\n", height(root))
	fmt.Println("Search 5:", Search(root, 5), "Search 99:", Search(root, 99))

	// 测试 2: 乱序插入
	root = nil
	for _, v := range []int{5, 2, 8, 1, 9, 3, 7, 6, 4} {
		root = Insert(root, v)
	}
	arr = nil
	Inorder(root, &arr)
	fmt.Println("\n乱序插入中序遍历:", arr)
	fmt.Printf("树高 = %d\n", height(root))

	// 测试 3: 删除
	root = Delete(root, 5)
	arr = nil
	Inorder(root, &arr)
	fmt.Println("删 5 后中序遍历:", arr)
	fmt.Println("Search 5:", Search(root, 5))
	fmt.Printf("树高 = %d\n", height(root))

	// 测试 4: 验证所有 BF ∈ {-1,0,1}
	if checkAllBF(root) {
		fmt.Println("\n所有节点 |BF| ≤ 1 ✓")
	}
}

// 验证整棵树每个节点的 |BF| ≤ 1
func checkAllBF(root *Node) bool {
	var walk func(*Node) bool
	walk = func(n *Node) bool {
		if n == nil {
			return true
		}
		bf := balanceFactor(n)
		if bf < -1 || bf > 1 {
			fmt.Printf("违反! 节点 %d 的 BF = %d\n", n.Val, bf)
			return false
		}
		return walk(n.Left) && walk(n.Right)
	}
	return walk(root)
}

func main() {
	runAVLDemo()
}