BST(1) - AVL树

Algorithm

January 17, 2023

前文介绍了二叉搜索树基础，这里是第一个自平衡的二叉搜索树 – AVL树，它是在 1962 年由前苏联的科学家 Georgy Adelson-Velsky 和 Evgenii Landis 提出的。它有最严格的平衡限制，因此有理论最佳的搜索效率。

基本概念

AVL树对树平衡的要求是（每个节点）左右树高的差距不超过 $1$ 。

定义 $\texttt{BF}$（Balance Factor）= 右子树高 - 左子树高，于是有 $\texttt{BF} \in $ { -1, 0, 1 } 。

于是我们需要一个额外的节点字段 $\texttt{height}$ 来辅助记录。

基础定义

共同结构

impl_node!();
impl_node_!({ height: i32 });
impl_tree!(AVL {});

impl_rotate_cleanup!(AVL ->
    fn rotate_cleanup(&self, x: Node<K, V>, z: Node<K, V>) {
        /* update height */
        x.update_height();
        z.update_height();
    }
);

impl<K: Debug, V> Debug for Node<K, V> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        if self.is_some() {
            write!(f, "{:?}(h: {})", key!(self), height!(self))
        }
        else {
            write!(f, "nil")
        }
    }
}

专属辅助

impl<K, V> Node<K, V> {
    fn update_height(&self) {
        if self.is_some() {
            height!(
                self,
                1 + max(
                    left!(self).height(),
                    right!(self).height()
                )
            );
        }
    }

    ////////////////////////////////////////////////////////////////////////////
    /// Static Stats

    fn height(&self) -> i32 {
        if self.is_none() {
            0
        }
        else {
            height!(self)
        }
    }

    fn bf(&self) -> i32 {
        if self.is_none() {
            0
        }
        else {
            right!(self).height() - left!(self).height()
        }
    }
}

平衡校验

impl_balance_validation!(AVL ->
    #[cfg(test)]
    fn balance_validation(&mut self) {
        self.root.recalc_height();
        self.root.validate_bf();
    }
);


impl<K, V> Node<K, V> {
    /// Recursively validate BF:
    ///
    /// BF(X): H(right(X)) - H(left(X))
    ///
    /// BF(X) in {-1, 0, 1}
    ///
    #[cfg(test)]
    fn validate_bf(&self) {
        assert!(
            self.bf().abs() < 2
        );

        if self.is_some() {
            left!(self).validate_bf();
            right!(self).validate_bf();
        }
    }

    /// Recursively calculate height stats dynamically instead of using static height
    #[cfg(test)]
    fn recalc_height(&self) {
        if self.is_some() {
            left!(self).recalc_height();
            right!(self).recalc_height();

            self.update_height();
        }
    }
}

重平衡

一般的重平衡操作

从入口节点开始，一路向上直到根节点，检查每个节点是否失衡，如果失衡根据前文讲的情况：

较高的子树的方向与子树较高子树的方向一致，一次单旋；
方向不一致，双旋

// impl<K: Ord, V> AVL<K, V> 
// ...

/// Bottom up fixing
fn retracing(&mut self, ent: Node<K, V>)
{
    let mut p = ent;

    while p.is_some() {
        p.update_height();

        if p.bf().abs() > 1 {
            let high =
            if right!(p).height() > left!(p).height() {
                Right
            }
            else {
                Left
            };

            let z = child!(p, high);

            if child!(z, high).height() >= child!(z, high.rev()).height() {
                p = rotate!(self, p, high.rev());
            }
            else {
                p = double_rotate!(self, p, high.rev());
            }
        }

        p = paren!(p).upgrade();
    }

}

// ...

对于插入操作简化后的重平衡：

对于插入操作的重平衡操作，可以简化检查较高子树的方向这一步，因为总是检查重平衡所在的节点较高：

// impl ...

#[allow(unused)]
fn insert_retracing(&mut self, mut y: Node<K, V>)
{
    /* x
       |
       z
       |
       y
    */

    let mut z = paren!(y).upgrade();

    while z.is_some() {
        z.update_height();

        let x = paren!(z).upgrade();

        if x.bf().abs() > 1 {
            let index_of_z = index_of_child!(x, z);
            let index_of_y = index_of_child!(z, y);

            if index_of_z == index_of_y {
                z = rotate!(self, x, index_of_z.rev());
            }
            else {
                z = double_rotate!(self, x, index_of_z.rev());
            }
        }

        y = z;
        z = paren!(y).upgrade();
    }

}

// ...

总装

impl<K: Ord, V> AVL<K, V> {

    ////////////////////////////////////////////////////////////////////////////
    /// Public API

    pub fn new() -> Self {
        Self {
            root: Node::none()
        }
    }


    pub fn insert(&mut self, k: K, v: V) -> Option<V>
    {
        let z = node!( BST { k, v, height: 1 });

        let popped = bst_insert!(self, z.clone());

        // self.insert_retracing(z);
        self.retracing(z);

        popped
    }


    pub fn remove<Q>(&mut self, k: &Q) -> Option<V>
    where K: Borrow<Q> + Debug, Q: Ord + ?Sized, V: Debug
    {

        let z = bst_search!(self.root, k);

        if z.is_none() {
            None
        }
        else {
            let retracing_entry = bst_delete!(self, z);
            self.retracing(retracing_entry);

            Some(unboxptr!(unwrap_into!(z).val))
        }
    }
}