数据结构学习-AVL树

在前面的 二叉搜索树 中我们已经知道对于已经排序好的数据它能够很好的提高数据的查找，添加，删除操作的性能，时间复杂度都是O(log n)。

但是如果二叉搜索树是下面这样：

如果下一个的插入节点都是插入到最底部的右节点，比如插入的是 10，那么它的插入时间复杂度其实是 O(n)。所以这就涉及二叉搜索树的一个平衡问题。

介绍

Georgy Adelson-Velsky 和 Evgenii Landis 提出的自平衡二叉搜索树，也就是 AVL 树（他们名字的首字母合在一起的命名）。它能保证数据插入，删除等操作的时间复杂度为 O(log n)。

优化不平衡的二叉搜索树的关键是让它变成平衡的二叉树，比如下面的情况：

(图1)

(图2)

(图1) (图2) 都是平衡二叉树，除了最底层可以不填满，每一层的节点都已经填满。其中(图1)是完美平衡二叉树。

因此下面的二叉树就不是平衡二叉树：

（7节点处少了左节点）

实现

因为AVL树是在二叉搜索树的基础上优化成平衡二叉树的，所以下面的代码是基于二次搜索树实现的。代码中把 BinaryNode 重命名为 AVLNode， BinarySearchTree 重命名为 AVLTree。

平衡因子

为了判断是否平衡，首先需要一个平衡因子的东西。平衡因子的值为左节点的高度减去右节点的高度，下面是实现代码：

class AVLNode<Element> {
    
    ...
    ...

    var height = 0

    // 平衡因子
    var balanceFactor: Int {
        leftHeight - rightHeight
    }
    
    // 左节点树的高度
    var leftHeight: Int {
        leftChild?.height ?? -1
    }
    
    // 右节点树的高度
    var rightHeight: Int {
        rightChild?.height ?? -1
    }
}

比如图片中的 节点1，它的左节点为空，所以左边高度为-1，右节点高度为0，因此它的平衡因子 = -1 。

旋转

为什么要旋转？这是因为插入节点以后，如果这时候的二叉树不是平衡二叉树，那么就需要把它调整为平衡二叉树，而调整的关键就是旋转某些节点。

左旋转

如下图所示的调整 节点10 使二叉树平衡：

代码实现：

extension AVLTree {

    private func leftRotate(_ node: AVLNode<Element>) -> AVLNode<Element> {
        let pivot = node.rightChild!
        node.rightChild = pivot.leftChild
        pivot.leftChild = node
        
        node.height = max(node.leftHeight, node.rightHeight) + 1
        pivot.height = max(pivot.leftHeight, pivot.rightHeight) + 1
        
        return pivot
    }
}

右旋转

代码实现：

extension AVLTree {
    private func rightRotate(_ node: AVLNode<Element>) -> AVLNode<Element> {
        let pivot = node.leftChild!
        node.leftChild = pivot.rightChild
        pivot.rightChild = node
        
        node.height = max(node.leftHeight, node.rightHeight) + 1
        pivot.height = max(pivot.leftHeight, pivot.rightHeight) + 1
        
        return pivot
    }
}

右左旋转

代码实现：

extension AVLTree {
    private func rightLeftRotate(_ node: AVLNode<Element>) -> AVLNode<Element> {
        guard let rightChild = node.rightChild else {
            return node
        }
        node.rightChild = rightRotate(rightChild)
        return leftRotate(node)
    }
}

左右旋转

代码实现：

extension AVLTree {
    private func leftRightRotate(_ node: AVLNode<Element>) -> AVLNode<Element> {
        guard let leftChild = node.leftChild else {
            return node
        }
        node.leftChild = leftRotate(leftChild)
        return rightRotate(node)
    }
}

平衡

保持平衡的实现关键点：

• 平衡因子等于2的时候，表明左节点的高度大于右节点的高度；

• 平衡因子等于-2的时候，表明右节点的高度大于左节点的高度。

代码实现：

extension AVLTree {
    private func balanced(_ node: AVLNode<Element>) -> AVLNode<Element> {
        switch node.balanceFactor {
        case 2:
            if let leftChild = node.leftChild, leftChild.balanceFactor == -1 {
                return leftRightRotate(node)
            } else {
                return rightRotate(node)
            }
        case -2:
            if let rightChild = node.rightChild, rightChild.balanceFactor == 1 {
                return rightLeftRotate(node)
            } else {
                return leftRotate(node)
            }
        default:
            return node
        }
    }
}

插入节点

实现代码：

extension AVLTree {
    
    mutating func insert(_ value: Element) {
        root = insert(from: root, value: value)
    }
    
    private func insert(from node: AVLNode<Element>?, value: Element) -> AVLNode<Element> {
        guard let node = node else {
            return AVLNode(value: value)
        }
        if value < node.value {
            node.leftChild = insert(from: node.leftChild, value: value)
        } else {
            node.rightChild = insert(from: node.rightChild, value: value)
        }
        let balancedNode = balanced(node)
        balancedNode.height = max(balancedNode.leftHeight, balancedNode.rightHeight) + 1
        return balancedNode
    }
    
}

测试：

func testInsert() {
    var tree = AVLTree<Int>()
    for i in 0..<12 {
        tree.insert(i)
    }
    print(tree)
}

输出：

  ┌──11
 ┌──10
 │ └──nil
┌──9
│ └──8
7
│  ┌──6
│ ┌──5
│ │ └──4
└──3
 │ ┌──2
 └──1
  └──0

在实际的测试中可以一步步的插入节点测试，查看每一步的插入节点后的输出是怎样的，这样能够对里面的步骤有更清晰的认识。

删除节点

实现代码：

extension AVLTree {
    
    mutating func remove(_ value: Element) {
        root = remove(node: root, value: value)
    }
    
    private func remove(node: AVLNode<Element>?, value: Element) -> AVLNode<Element>? {
        guard let node = node else {
            return nil
        }
        
        if value == node.value {
            // 叶子节点
            if node.leftChild == nil && node.rightChild == nil {
                return nil
            }
            // 只有一个右节点
            if node.leftChild == nil {
                return node.rightChild
            }
            // 只有一个左节点
            if node.rightChild == nil {
                return node.leftChild
            }
            // 有两个子节点，替换位置，删除替换后右节点的树中的最小节点
            node.value = node.rightChild!.min.value
            node.rightChild = remove(node: node.rightChild, value: node.value)
        } else if value < node.value  {
            node.leftChild = remove(node: node.leftChild, value: value)
        } else {
            node.rightChild = remove(node: node.rightChild, value: value)
        }
        
        // avl 的修改处
        let balancedNode = balanced(node)
        balancedNode.height = max(balancedNode.leftHeight, balancedNode.rightHeight) + 1
        return balancedNode
    }
}

测试：

func testRemove() {
    var tree = AVLTree<Int>()
    for i in 0..<12 {
        tree.insert(i)
    }
    
    print("after remove 3, 10")
    tree.remove(3)
    tree.remove(10)
    
    print(tree)
}

输出：

 ┌──11
┌──9
│ └──8
7
│  ┌──6
│ ┌──5
│ │ └──nil
└──4
 │ ┌──2
 └──1
  └──0

每一次的插入和删除节点后的结果都是一棵平衡二叉树。