（22）Go實(shí)現(xiàn)AVL樹-實(shí)現(xiàn)和測試

繼上一篇《（21）Go實(shí)現(xiàn)AVL樹-算法解析》的后續(xù)
http://www.reibang.com/p/943243f5ee1e

AVL樹的定義和實(shí)現(xiàn)
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
}

// 添加/更新節(jié)點(diǎn)
func (tree *avlTree) Add(key, val int) {
    isAdd, nd := tree.Root.add(key, val)
    tree.size += isAdd
    tree.Root = nd
}

// 遞歸寫法:向樹的root節(jié)點(diǎn)中插入key,val
// 返回1,代表加了節(jié)點(diǎn)
// 返回0,代表沒有添加新節(jié)點(diǎn),只更新key對應(yīng)的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值更新,節(jié)點(diǎn)數(shù)量不增加,isAdd = 0
        nd.value = val
    }

    // 更新節(jié)點(diǎn)高度和維護(hù)平衡
    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為根節(jié)點(diǎn)中key節(jié)點(diǎn),返回更新了高度和維持平衡的新nd節(jié)點(diǎn)
// 返回值int == 1 表明有節(jié)點(diǎn)被刪除,0 表明沒有節(jié)點(diǎn)被刪除
func (nd *node) remove(key int) (int, *node) {
    // 找不到key對應(yīng)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: // 待刪除節(jié)點(diǎn)左子樹為空的情況
            retNode = nd.right
        case nd.right == nil: // 待刪除節(jié)點(diǎn)右子樹為空的情況
            retNode = nd.left
        default:
            // 待刪除節(jié)點(diǎn)左右子樹均不為空的情況
            // 找到比待刪除節(jié)點(diǎn)大的最小節(jié)點(diǎn),即右子樹的最小節(jié)點(diǎn)
            retNode := nd.right.getMinNode()
            // TODO: 這步好好理解,維護(hù)平衡性
            _, retNode.right = nd.right.remove(key)

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

    // 前面刪除節(jié)點(diǎn)后,返回retNode有可能為空,這樣在執(zhí)行下面的語句會產(chǎn)生異常，加這步判定避免異常
    if retNode == nil {
        return isRemove, retNode
    }

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

// 對節(jié)點(diǎn)y進(jìn)行向右旋轉(zhuǎn)操作砸民，返回旋轉(zhuǎn)后新的根節(jié)點(diǎn)x
//        y                              x
//       / \                           /   \
//      x   T4     向右旋轉(zhuǎn) (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值: 旋轉(zhuǎn)后只有x,y的子樹發(fā)生變化,所以只更新x,y的height值
    // x的height值和y的height值相關(guān),先更新y,再更新x
    y.height = 1 + max(y.left.getHeight(), y.right.getHeight())
    x.height = 1 + max(x.left.getHeight(), x.right.getHeight())
    return x
}

// 對節(jié)點(diǎn)y進(jìn)行向左旋轉(zhuǎn)操作探熔，返回旋轉(zhuǎn)后新的根節(jié)點(diǎn)x
//    y                             x
//  /  \                          /   \
// T1   x      向左旋轉(zhuǎn) (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節(jié)點(diǎn)
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是否仍然是一顆二分搜索樹
// 思路: 二分搜索數(shù)如果用中序遍歷時(shí),所有元素都是從小到大排列
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是否是一顆平衡二叉樹(每個節(jié)點(diǎn)的左右子樹高度差不能超過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
    }
    // 能到這里說明:當(dāng)前該節(jié)點(diǎn)滿足平衡二叉樹條件
    // 再去判斷該節(jié)點(diǎn)的左右子樹是否也是平衡二叉樹
    return nd.left.isBalanced() && nd.right.isBalanced()
}

// 找出以nd為根節(jié)點(diǎn)中最小值的節(jié)點(diǎn)
func (nd *node) getMinNode() *node {
    if nd.left == nil {
        return nd
    }
    return nd.left.getMinNode()
}

// 更新節(jié)點(diǎn)高度和維護(hù)平衡
func (nd *node) updateHeightAndBalance(isChange int) *node {
    // 0說明無改變,不必更新
    if isChange == 0 {
        return nd
    }

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

    // 平衡維護(hù)
    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("測試結(jié)束")
    fmt.Println("is BST:", a.IsBST())
    fmt.Println("is Balanced:", a.IsBalanced())
}

測試結(jié)果
[]
----
測試結(jié)束
is BST: true
is Balanced: true

有bug歡迎指出