繼上一篇 《（21）Go實現AVL樹-算法解析》 的后續
http://www.lxweimin.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歡迎指出
