（22）Go实现AVL树-实现和测试

继上一篇《（21）Go实现AVL树-算法解析》的后续
https://www.jianshu.com/p/943243f5ee1e

AVL树的定义和实现
type node struct {
    key    int
    value  int
    height int
    left   *node
    right  *node
}

type avlTree struct {
    size int
    Root *node
}

func NewAvlTree() *avlTree {
    return new(avlTree)
}

func (tree *avlTree) GetSize() int {
    return tree.size
}

// 获取高度
func (nd *node) getHeight() int {
    if nd == nil {
        return 0
    }
    return nd.height
}

// 获取的平衡因子
func (nd *node) getBalanceFactor() int {
    if nd == nil {
        return 0
    }
    return nd.left.getHeight() - nd.right.getHeight()
}

func max(int1, int2 int) int {
    if int1 >= int2 {
        return int1
    }
    return int2
}

// 添加/更新节点
func (tree *avlTree) Add(key, val int) {
    isAdd, nd := tree.Root.add(key, val)
    tree.size += isAdd
    tree.Root = nd
}

// 递归写法:向树的root节点中插入key,val
// 返回1,代表加了节点
// 返回0,代表没有添加新节点,只更新key对应的value值
func (nd *node) add(key, val int) (int, *node) {
    if nd == nil {
        return 1, &node{key, val, 1, nil, nil}
    }

    isAdd := 0
    if key < nd.key {
        isAdd, nd.left = nd.left.add(key, val)
    } else if key > nd.key {
        isAdd, nd.right = nd.right.add(key, val)
    } else { // nd.key == key
        // 对value值更新,节点数量不增加,isAdd = 0
        nd.value = val
    }

    // 更新节点高度和维护平衡
    nd = nd.updateHeightAndBalance(isAdd)
    return isAdd, nd
}

func (tree *avlTree) Remove(key int) error {
    if tree.Root == nil {
        return errors.New(
            "failed to remove,avlTree is empty.")
    }
    var isRemove int
    isRemove, tree.Root = tree.Root.remove(key)
    tree.size -= isRemove
    return nil
}

// 删除nd为根节点中key节点,返回更新了高度和维持平衡的新nd节点
// 返回值int == 1 表明有节点被删除,0 表明没有节点被删除
func (nd *node) remove(key int) (int, *node) {
    // 找不到key对应node,返回0,nil
    if nd == nil {
        return 0, nil
    }

    var retNode *node
    var isRemove int
    switch {
    case key < nd.key:
        isRemove, nd.left = nd.left.remove(key)
        retNode = nd
    case key > nd.key:
        isRemove, nd.right = nd.right.remove(key)
        retNode = nd
    default:
        switch {
        case nd.left == nil: // 待删除节点左子树为空的情况
            retNode = nd.right
        case nd.right == nil: // 待删除节点右子树为空的情况
            retNode = nd.left
        default:
            // 待删除节点左右子树均不为空的情况
            // 找到比待删除节点大的最小节点,即右子树的最小节点
            retNode := nd.right.getMinNode()
            // TODO: 这步好好理解,维护平衡性
            _, retNode.right = nd.right.remove(key)

            retNode.left = nd.left
        }
        isRemove = 1
    }

    // 前面删除节点后,返回retNode有可能为空,这样在执行下面的语句会产生异常，加这步判定避免异常
    if retNode == nil {
        return isRemove, retNode
    }

    retNode = retNode.updateHeightAndBalance(isRemove)
    return isRemove, retNode
}

// 对节点y进行向右旋转操作，返回旋转后新的根节点x
//        y                              x
//       / \                           /   \
//      x   T4     向右旋转 (y)        z     y
//     / \       - - - - - - - ->    / \   / \
//    z   T3                       T1  T2 T3 T4
//   / \
// T1   T2
func (y *node) rightRotate() *node {
    x := y.left
    y.left = x.right
    x.right = y

    // 更新height值: 旋转后只有x,y的子树发生变化,所以只更新x,y的height值
    // x的height值和y的height值相关,先更新y,再更新x
    y.height = 1 + max(y.left.getHeight(), y.right.getHeight())
    x.height = 1 + max(x.left.getHeight(), x.right.getHeight())
    return x
}

// 对节点y进行向左旋转操作，返回旋转后新的根节点x
//    y                             x
//  /  \                          /   \
// T1   x      向左旋转 (y)       y     z
//     / \   - - - - - - - ->   / \   / \
//   T2  z                     T1 T2 T3 T4
//      / \
//     T3 T4
func (y *node) leftRotate() *node {
    x := y.right
    y.right = x.left
    x.left = y

    // 更新height值
    y.height = 1 + max(y.left.getHeight(), y.right.getHeight())
    x.height = 1 + max(x.left.getHeight(), x.right.getHeight())
    return x
}

func (tree *avlTree) Contains(key int) bool {
    return tree.Root.contains(key)
}

// 以root为根的树中是否包含key节点
func (nd *node) contains(key int) bool {
    if nd == nil {
        return false
    }

    if nd.key == key {
        return true
    }

    if key < nd.key {
        return nd.left.contains(key)
    }
    return nd.right.contains(key)
}

// 中序遍历打印出key,val,height
func (tree *avlTree) PrintInOrder() {
    resp := [][]int{}
    tree.Root.printInOrder(&resp)
    fmt.Println(resp)
}

func (nd *node) printInOrder(resp *[][]int) {
    if nd == nil {
        return
    }
    nd.left.printInOrder(resp)
    *resp = append(*resp, []int{nd.key, nd.value, nd.height})
    nd.right.printInOrder(resp)
}

// 中序遍历所有key,用来辅助判断是否是二分搜索树
func (nd *node) traverseInOrderKey(resp *[]int) {
    if nd == nil {
        return
    }

    nd.left.traverseInOrderKey(resp)
    *resp = append(*resp, nd.key)
    nd.right.traverseInOrderKey(resp)
}

// 判断avlTree是否仍然是一颗二分搜索树
// 思路: 二分搜索数如果用中序遍历时,所有元素都是从小到大排列
func (tree *avlTree) IsBST() bool {
    buf := []int{}
    tree.Root.traverseInOrderKey(&buf)
    length := len(buf)
    for i := 1; i < length; i++ {
        if buf[i-1] > buf[i] {
            return false
        }
    }
    return true
}

// 判断avlTree是否是一颗平衡二叉树(每个节点的左右子树高度差不能超过1)
func (tree *avlTree) IsBalanced() bool {
    return tree.Root.isBalanced()
}

func (nd *node) isBalanced() bool {
    if nd == nil {
        return true
    }

    if nd.getBalanceFactor() > 1 || nd.getBalanceFactor() < int(-1) {
        return false
    }
    // 能到这里说明:当前该节点满足平衡二叉树条件
    // 再去判断该节点的左右子树是否也是平衡二叉树
    return nd.left.isBalanced() && nd.right.isBalanced()
}

// 找出以nd为根节点中最小值的节点
func (nd *node) getMinNode() *node {
    if nd.left == nil {
        return nd
    }
    return nd.left.getMinNode()
}

// 更新节点高度和维护平衡
func (nd *node) updateHeightAndBalance(isChange int) *node {
    // 0说明无改变,不必更新
    if isChange == 0 {
        return nd
    }

    // 更新高度
    nd.height = 1 + max(nd.left.getHeight(),
        nd.right.getHeight())

    // 平衡维护
    if nd.getBalanceFactor() > 1 {
        // 左左LL
        if nd.left.getBalanceFactor() >= 0 {
            nd = nd.rightRotate()
        } else { // 左右LR
            nd.left = nd.left.leftRotate()
            nd = nd.rightRotate()
        }
        return nd
    }

    if nd.getBalanceFactor() < int(-1) {
        // 右右RR
        if nd.right.getBalanceFactor() <= 0 {
            nd = nd.leftRotate()
        } else { // 右左RL
            nd.right = nd.right.rightRotate()
            nd = nd.leftRotate()
        }
    }
    return nd
}

测试代码
    a := avltree1.NewAvlTree()
    nums := []int{}
    a.PrintInOrder()
    fmt.Println("----")
    for i := 0; i < 100000; i++ {
        a.Add(100000, 0)
        if !a.IsBST() || !a.IsBalanced() {
            fmt.Println("代码有错误", a.IsBST(), a.IsBalanced())
            break
        }
    }

    for i := 0; i < 200000; i++ {
        nums = append(nums, rand.Intn(200000))
    }
    for _, v := range nums {
        a.Remove(v)
        if !a.IsBST() || !a.IsBalanced() {
            fmt.Println("代码有错误", a.IsBST(), a.IsBalanced())
            break
        }
    }

    fmt.Println("测试结束")
    fmt.Println("is BST:", a.IsBST())
    fmt.Println("is Balanced:", a.IsBalanced())
}

测试结果
[]
----
测试结束
is BST: true
is Balanced: true

有bug欢迎指出