平衡二叉查找树

平衡二叉查找树
- 平衡因子
添加节点
完整例子

平衡二叉搜索树（Self-balancing binary search tree）又被称为AVL树（有别于AVL算法），且具有以下性质：它是一棵空树或它的左右两个子树的高度差的绝对值不超过1，并且左右两个子树都是一棵平衡二叉树。

平衡因子

某结点的左子树与右子树的高度(深度)差即为该结点的平衡因子（BF,Balance Factor）。平衡二叉树上所有结点的平衡因子只可能是 -1，0 或 1。如果某一结点的平衡因子绝对值大于1则说明此树不是平衡二叉树。为了方便计算每一结点的平衡因子我们可以为每个节点赋予height这一属性，表示此节点的高度。

type avlTree struct {
    value, height       int
    left, right, parent *avlTree
}
func NewAvlTree(value int) *avlTree {
    return &avlTree{
        value:  value,
        height: 1,
    }
}
// 得到树高
func (a *avlTree) getHeight() int {
    if a == nil {
        return 0
    }
    return a.height
}
// 获取节点的平衡因子
func (a *avlTree) getBalanceFactor() int {
    if a == nil {
        return 0
    }
    return a.left.getHeight() - a.right.getHeight()
}

添加节点

往平衡二叉树中添加节点很可能会导致二叉树失去平衡，所以我们需要在每次插入节点后进行平衡的维护操作。插入节点破坏平衡性有如下四种情况：

右旋

意思是向左子树（L）的左孩子（L）中插入新节点后导致不平衡，这种情况下需要右旋操作

我们将这种情况抽象出来，得到下图：

我们需要对节点y进行平衡的维护。步骤如下图所示：

func (a *avlTree) rightRotate() *avlTree {
    x := a.left
    a.left = nil
    x.parent = a.parent
    a.parent = x
    xr := x.right
    x.right = a
    a.left = xr
    if xr != nil {
        xr.parent = a
    }
    x.height = 1 + int(math.Max(float64(x.left.getHeight()), float64(x.right.getHeight())))
    a.height = 1 + int(math.Max(float64(a.left.getHeight()), float64(a.right.getHeight())))
    return x
}

左旋

意思是向右子树（R）的右孩子（R）中插入新节点后导致不平衡，这种情况下需要左旋操作

我们将这种情况抽象出来，得到下图：

我们需要对节点y进行平衡的维护。步骤如下图所示：

func (a *avlTree) leftRotate() *avlTree {
    x := a.right
    a.right = nil
    x.parent = a.parent
    a.parent = x
    xl := x.left
    x.left = a
    a.right = xl
    if xl != nil {
        xl.parent = a
    }
    x.height = 1 + int(math.Max(float64(x.left.getHeight()), float64(x.right.getHeight())))
    a.height = 1 + int(math.Max(float64(a.left.getHeight()), float64(a.right.getHeight())))
    return x
}

左旋再右旋

意思是向左子树（L）的右孩子（R）中插入新节点后导致不平衡，这种情况下需要左旋操作

我们将这种情况抽象出来，得到下图：

我们需要对节点y进行平衡的维护。步骤如下图所示：

右旋再左旋

意思是向右子树（R）的左孩子（L）中插入新节点后导致不平衡，这种情况下需要左旋操作

我们将这种情况抽象出来，得到下图：

我们需要对节点y进行平衡的维护。步骤如下图所示：

什么情况用那种旋转

左树比右树高2或以上，且左子树的左子叶节点高度大于右子叶节点

此时使用右旋

右树比左树高2或以上，且右子树的右子叶节点高度左大于子叶节点

此时使用左旋

左树比右树高2或以上，且左子树的左子叶节点高度小于右子叶节点

此时使用左旋再右旋

右树比左树高2或以上，且右子树的左子叶节点高度大于右子叶节点

此时使用右旋再左旋

添加节点代码

func (a *avlTree) Insert(value int) *avlTree {
    if a == nil {
        return NewAvlTree(value)
    }
    v := a.compareTo(value)
    if v < 0 {
        a.left = a.left.Insert(value)
        if a.left.parent == nil {
            a.left.parent = a
        }
    } else {
        a.right = a.right.Insert(value)
        if a.right.parent == nil {
            a.right.parent = a
        }
    }
    //更新height
    a.height = 1 + int(math.Max(float64(a.left.getHeight()), float64(a.right.getHeight())))
    balanceFactor := a.getBalanceFactor()
    // 右旋
    if balanceFactor > 1 && a.left.getBalanceFactor() > 0 {
        return a.rightRotate()
    }
    // 左旋
    if balanceFactor < -1 && a.right.getBalanceFactor() < 0 {
        return a.leftRotate()
    }
    // 先左旋，再右旋
    if balanceFactor > 1 && a.left.getBalanceFactor() < 0 {
        a.left = a.left.leftRotate()
        return a.rightRotate()
    }
    // 先右旋，再左旋
    if balanceFactor < -1 && a.right.getBalanceFactor() > 0 {
        a.right = a.right.rightRotate()
        return a.leftRotate()
    }

    return a
}

删除节点

删除节点与二叉查询书一样，不过删除后需要考虑平衡的问题

func (a *avlTree) Remove(value int) {
    v := a.compareTo(value)
    if v < 0 {
        a.left.Remove(value)
    } else if v > 0 {
        a.right.Remove(value)
    } else if a.left != nil && a.right != nil {
        // 删除的节点同时存在左右节点
        // 找个后继节点
        successor := a.right.Min()
        // 将删除节点的左节点的父节点指向后继节点
        a.left.parent = successor
        // 将后继节点的左节点指向删除节点的左节点
        successor.left = a.left
        // 如果后继节点不是一个孤立节点
        if successor.parent != a {
            // 将后继节点的右节点指向删除节点的右节点
            sr := successor.right
            successor.right = a.right
            a.right.parent = successor
            // 将后继节点的父节点的左节点指向successor的右节点
            if sr != nil {
                sr.parent = successor.parent
            }
            successor.parent.left = sr
        }
        // 将后继节点的父节点指向删除节点的父节点
        successor.parent = a.parent
        // 将后继节点关联到树里
        if a.parent == nil {
            *a = *successor
            return
        }
        if a.parent.left == a {
            a.parent.left = successor
        } else {
            a.parent.right = successor
        }
    } else if a.left != nil || a.right != nil {
        // 删除的节点存在左节点或右节点
        var child *avlTree
        if a.left != nil {
            child = a.left
        } else {
            child = a.right
        }
        child.parent = a.parent
        if a.parent == nil {
            *a = *child
            return
        }
        if a.parent.left == a {
            a.parent.left = child
        } else {
            a.parent.right = child
        }

    } else {
        //  删除的节点没左右节点
        if a.parent == nil {
            *a = avlTree{}
            return
        }
        if a.parent.left == a {
            a.parent.left = nil
        } else {
            a.parent.right = nil
        }
    }

    //更新height
    a.height = 1 + int(math.Max(float64(a.left.getHeight()), float64(a.right.getHeight())))
    balanceFactor := a.getBalanceFactor()
    // 右旋
    if balanceFactor > 1 && a.left.getBalanceFactor() > 0 {
        a.rightRotate()
    }
    // 左旋
    if balanceFactor < -1 && a.right.getBalanceFactor() < 0 {
        a.leftRotate()
    }
    // 先左旋，再右旋
    if balanceFactor > 1 && a.left.getBalanceFactor() < 0 {
        a.left = a.left.leftRotate()
        a.rightRotate()
    }
    // 先右旋，再左旋
    if balanceFactor < -1 && a.right.getBalanceFactor() > 0 {
        a.right = a.right.rightRotate()
        a.leftRotate()
    }
}

完整例子

package main

import (
    "fmt"
    "math"
)

type avlTree struct {
    value, height       int
    left, right, parent *avlTree
}

func NewAvlTree(value int) *avlTree {
    return &avlTree{
        value:  value,
        height: 1,
    }
}

// 得到树高
func (a *avlTree) getHeight() int {
    if a == nil {
        return 0
    }
    return a.height
}

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

func (a *avlTree) rightRotate() *avlTree {
    x := a.left
    a.left = nil
    x.parent = a.parent
    a.parent = x
    xr := x.right
    x.right = a
    a.left = xr
    if xr != nil {
        xr.parent = a
    }
    x.height = 1 + int(math.Max(float64(x.left.getHeight()), float64(x.right.getHeight())))
    a.height = 1 + int(math.Max(float64(a.left.getHeight()), float64(a.right.getHeight())))
    return x
}

func (a *avlTree) leftRotate() *avlTree {
    x := a.right
    a.right = nil
    x.parent = a.parent
    a.parent = x
    xl := x.left
    x.left = a
    a.right = xl
    if xl != nil {
        xl.parent = a
    }
    x.height = 1 + int(math.Max(float64(x.left.getHeight()), float64(x.right.getHeight())))
    a.height = 1 + int(math.Max(float64(a.left.getHeight()), float64(a.right.getHeight())))
    return x
}

func (a *avlTree) compareTo(value int) int {
    return value - a.value
}

func (a *avlTree) Insert(value int) *avlTree {
    if a == nil {
        return NewAvlTree(value)
    }
    v := a.compareTo(value)
    if v < 0 {
        a.left = a.left.Insert(value)
        if a.left.parent == nil {
            a.left.parent = a
        }
    } else {
        a.right = a.right.Insert(value)
        if a.right.parent == nil {
            a.right.parent = a
        }
    }
    //更新height
    a.height = 1 + int(math.Max(float64(a.left.getHeight()), float64(a.right.getHeight())))
    balanceFactor := a.getBalanceFactor()
    // 右旋
    if balanceFactor > 1 && a.left.getBalanceFactor() > 0 {
        return a.rightRotate()
    }
    // 左旋
    if balanceFactor < -1 && a.right.getBalanceFactor() < 0 {
        return a.leftRotate()
    }
    // 先左旋，再右旋
    if balanceFactor > 1 && a.left.getBalanceFactor() < 0 {
        a.left = a.left.leftRotate()
        return a.rightRotate()
    }
    // 先右旋，再左旋
    if balanceFactor < -1 && a.right.getBalanceFactor() > 0 {
        a.right = a.right.rightRotate()
        return a.leftRotate()
    }

    return a
}

// GetAll 获取树种所有的元素值（并按从小到大排序）
func (a *avlTree) GetAll() []int {
    var values []int
    return appendValues(values, a)
}

func appendValues(values []int, t *avlTree) []int {
    if t != nil {
        values = appendValues(values, t.left)
        values = append(values, t.value)
        values = appendValues(values, t.right)
    }
    return values
}

func (a *avlTree) Min() *avlTree {
    if a == nil {
        fmt.Println("tree is empty")
        return a
    }

    if a.left == nil {
        return a
    }

    return a.left.Min()
}

func (a *avlTree) Max() *avlTree {
    if a == nil {
        fmt.Println("tree is empty")
        return a
    }

    if a.right == nil {
        return a
    }

    return a.right.Max()
}

func (a *avlTree) Remove(value int) {
    v := a.compareTo(value)
    if v < 0 {
        a.left.Remove(value)
    } else if v > 0 {
        a.right.Remove(value)
    } else if a.left != nil && a.right != nil {
        // 删除的节点同时存在左右节点
        // 找个后继节点
        successor := a.right.Min()
        // 将删除节点的左节点的父节点指向后继节点
        a.left.parent = successor
        // 将后继节点的左节点指向删除节点的左节点
        successor.left = a.left
        // 如果后继节点不是一个孤立节点
        if successor.parent != a {
            // 将后继节点的右节点指向删除节点的右节点
            sr := successor.right
            successor.right = a.right
            a.right.parent = successor
            // 将后继节点的父节点的左节点指向successor的右节点
            if sr != nil {
                sr.parent = successor.parent
            }
            successor.parent.left = sr
        }
        // 将后继节点的父节点指向删除节点的父节点
        successor.parent = a.parent
        // 将后继节点关联到树里
        if a.parent == nil {
            *a = *successor
            return
        }
        if a.parent.left == a {
            a.parent.left = successor
        } else {
            a.parent.right = successor
        }
    } else if a.left != nil || a.right != nil {
        // 删除的节点存在左节点或右节点
        var child *avlTree
        if a.left != nil {
            child = a.left
        } else {
            child = a.right
        }
        child.parent = a.parent
        if a.parent == nil {
            *a = *child
            return
        }
        if a.parent.left == a {
            a.parent.left = child
        } else {
            a.parent.right = child
        }

    } else {
        //  删除的节点没左右节点
        if a.parent == nil {
            *a = avlTree{}
            return
        }
        if a.parent.left == a {
            a.parent.left = nil
        } else {
            a.parent.right = nil
        }
    }

    //更新height
    a.height = 1 + int(math.Max(float64(a.left.getHeight()), float64(a.right.getHeight())))
    balanceFactor := a.getBalanceFactor()
    // 右旋
    if balanceFactor > 1 && a.left.getBalanceFactor() > 0 {
        a.rightRotate()
    }
    // 左旋
    if balanceFactor < -1 && a.right.getBalanceFactor() < 0 {
        a.leftRotate()
    }
    // 先左旋，再右旋
    if balanceFactor > 1 && a.left.getBalanceFactor() < 0 {
        a.left = a.left.leftRotate()
        a.rightRotate()
    }
    // 先右旋，再左旋
    if balanceFactor < -1 && a.right.getBalanceFactor() > 0 {
        a.right = a.right.rightRotate()
        a.leftRotate()
    }
}

// Contains 是否包含节点
func (a *avlTree) Contains(value int) bool {
    if a == nil {
        return false
    }
    v := a.compareTo(value)
    if v < 0 {
        return a.left.Contains(value)
    } else if v > 0 {
        return a.right.Contains(value)
    } else {
        return true
    }
}

func GetAvlTree() *avlTree {
    var a = NewAvlTree(30)
    in := []int{80, 20, 35, 34, 32, 40, 70, 75, 100}
    for _, v := range in {
        a = a.Insert(v)
    }
    return a
}

func main() {
    a1 := GetAvlTree()
    fmt.Println(a1.GetAll())

    a1.Remove(40)
    fmt.Println(a1.GetAll())
    fmt.Println(a1.Contains(30))
}