Inaha

使用 Rust 抽象 AVL Tree 并实现 Iter/IntoIter Trait (NLR/LNR/LRN)

最近在折腾一些基础知识,想趁机写一些博客。希望能坚持下来吧。

AVL 即平衡二叉树,是一种比较基础的数据结构。为了保证搜索速度,每次插入/删除之前都需要自我平衡。 使用 Rust 实现递归数据结构并非易事,对于咱这种初学者来说和编译器斗智斗勇是不可缺少的一个环节(

基本结构体

#[derive(Debug)]
enum AvlTree<T> {
    Null,
    Tree(Box<AvlNode<T>>),
}

#[derive(Debug)]
struct AvlNode<T> {
    val: T,
    left: AvlTree<T>,
    right: AvlTree<T>,
    bfactor: i8, // 平衡因子
}

也可以这样写, Option<T>take(), as_deref() 的默认实现。

type AvlTree<T> = Option<Box<AvlNode<T>>>

AvlTree 实现基本操作

旋转

首先是基本的左旋与右旋。以左旋为例:
根节点 - R | 右子 - B | 假设存在的右子的左子 - A
- 将旧根节点(R)作为右子(B)的左子 (R),提升右子(B)为新的根节点。
- 若新根(B)已经有左子(A)与其冲突,那么使它成为新左子(R)的右子(A)

为保证平衡,还需要:
- 若子树需左旋使其平衡,先检查右子平衡因子。若它的右子是重的,先将它右旋
- 左旋子树

这里参考了
Problem-solving with algorithms and data structures using Rust - Shieber

impl<T> AvlTree<T>
where
    T: Ord,
{
    fn new() -> AvlTree<T> {
        Null
    }

    fn rebalance(&mut self) {
        match *self {
            Null => (),
            Tree(_) => match self.node().bfactor {
                // right
                -2 => {
                    let lbf = self.node().left.node().bfactor;
                    if lbf == -1 || lbf == 0 {
                        let (a, b) = if lbf == -1 { (0, 0) } else { (-1, 1) };
                        self.rotate_right();
                        self.node().right.node().bfactor = a;
                        self.node().bfactor = b;
                    } else if lbf == 1 {
                        let (a, b) = match self.node().left.node().right.node().bfactor {
                            -1 => (1, 0),
                            0 => (0, 0),
                            1 => (0, -1),
                            _ => unreachable!(),
                        };

                        self.node().left.rotate_left();
                        self.rotate_right();
                        self.node().right.node().bfactor = a;
                        self.node().left.node().bfactor = b;
                        self.node().bfactor = 0;
                    } else {
                        unreachable!()
                    }
                }
                // left
                2 => {
                    let rbf = self.node().right.node().bfactor;
                    if rbf == 1 || rbf == 0 {
                        let (a, b) = if rbf == 1 { (0, 0) } else { (1, -1) };
                        self.rotate_left();
                        self.node().left.node().bfactor = a;
                        self.node().bfactor = b;
                    } else if rbf == -1 {
                        let (a, b) = match self.node().right.node().left.node().bfactor {
                            1 => (-1, 0),
                            0 => (0, 0),
                            -1 => (0, 1),
                            _ => unreachable!(),
                        };

                        self.node().right.rotate_right();
                        self.rotate_left();
                        self.node().left.node().bfactor = a;
                        self.node().right.node().bfactor = b;
                        self.node().bfactor = 0;
                    } else {
                        unreachable!()
                    }
                }
                _ => (),
            },
        }
    }

    fn node(&mut self) -> &mut AvlNode<T> {
        match *self {
            Null => panic!("Empty tree"),
            Tree(ref mut n) => n,
        }
    }

    fn left_subtree(&mut self) -> &mut Self {
        match *self {
            Null => panic!("Empty tree"),
            Tree(ref mut node) => &mut node.left,
        }
    }

    fn right_subtree(&mut self) -> &mut Self {
        match *self {
            Null => panic!("Empty tree"),
            Tree(ref mut node) => &mut node.right,
        }
    }

    fn rotate_left(&mut self) {
        let mut v = self.take();
        let mut right = v.right_subtree().take();
        let right_left = right.left_subtree().take();
        *v.right_subtree() = right_left;
        *right.left_subtree() = v;
        *self = right;
    }

    fn rotate_right(&mut self) {
        let mut v = replace(self, Null);
        let mut left = replace(v.left_subtree(), Null);
        let left_right = replace(left.right_subtree(), Null);
        *v.left_subtree() = left_right;
        *left.right_subtree() = v;
        *self = left;
    }
}

插入

插入后需要重新平衡 AVL 树。

impl<T> AvlTree<T>
where
    T: Ord,
{
    fn insert(&mut self, val: T) -> (bool, bool) {
        let ret = match *self {
            Null => {
                let node = AvlNode {
                    val,
                    left: Null,
                    right: Null,
                    bfactor: 0,
                };
                *self = Tree(Box::new(node));

                (true, true)
            }
            Tree(ref mut node) => match node.val.cmp(&val) {
                Equal => (false, false),
                Less => {
                    let (inserted, deepened) = node.right.insert(val);

                    if deepened {
                        let ret = match node.bfactor {
                            -1 => (inserted, false),
                            0 => (inserted, true),
                            1 => (inserted, false),
                            _ => unreachable!(),
                        };

                        node.bfactor += 1;

                        ret
                    } else {
                        (inserted, deepened)
                    }
                }
                Greater => {
                    let (inserted, deepened) = node.left.insert(val);

                    if deepened {
                        let ret = match node.bfactor {
                            -1 => (inserted, false),
                            0 => (inserted, true),
                            1 => (inserted, false),
                            _ => unreachable!(),
                        };

                        node.bfactor -= 1;

                        ret
                    } else {
                        (inserted, deepened)
                    }
                }
            },
        };
        self.rebalance();

        ret
    }
}

inserted 代表是否已经插入, deepened 代表是否加深。
以插入右子为例:

Less => {
    let (inserted, deepened) = node.right.insert(val);

    if deepened {
        let ret = match node.bfactor {
            -1 => (inserted, false),
            0 => (inserted, true),
            1 => (inserted, false),
            _ => unreachable!(),
        };

        node.bfactor += 1;

        ret
    } else {
        (inserted, deepened)
    }
}

bfactor 的实际意义是 右侧深度 - 左侧深度,bfactor == 1 的时候不会加深,因为后续的再平衡操作会让树实际没有加深。

搜索

搜索的实现很简单。由于 AVL 已经十分有序,所以只需要查询子树根节点的值,大于它就去右子树找,小于它就去左子树找就好了。

impl<T> AvlTree<T>
where
    T: Ord,
{
    fn contains(&self, val: T) -> bool {
        match *self {
            Null => false,
            Tree(ref v) => match v.val.cmp(&val) {
                Equal => true,
                Less => v.right.contains(val),
                Greater => v.left.contains(val),
            },
        }
    }
}

AvlTree 实现 IntoIter/Iter/IterMut

这里还没写完(也是主要内容),先给出 NLR, LNR, LRN 的 IntoIter 实现与 LNR 的 Iter 实现。

一些需要实现的方法

impl<T> AvlTree<T>
where
    T: Ord,
{
    fn take(&mut self) -> AvlTree<T> {
        replace(self, AvlTree::Null)
    }

    fn len(&self) -> usize {
        match *self {
            Null => 0,
            Tree(ref v) => 1 + v.left.len() + v.right.len(),
        }
    }

    fn depth(&self) -> usize {
        match *self {
            Null => 0,
            Tree(ref v) => 1 + max(v.left.depth(), v.right.depth()),
        }
    }

    fn is_empty(&self) -> bool {
        matches!(*self, Null)
    }
}

IntoIter 的实现

首先是 IntoIter, NLR 与 LRN 的实现较为简单

struct AvlNLRIntoIter<T> {
    stack: Vec<AvlTree<T>>,
}
struct AvlLRNIntoIter<T> {
    stack: Vec<AvlTree<T>>,
}

impl<T> AvlTree<T>
where
    T: Ord,
{
    fn into_iter_nrl(self) -> AvlNLRIntoIter<T> {
        AvlNLRIntoIter { stack: vec![self] }
    }

    fn into_iter_rln(self) -> AvlLRNIntoIter<T> {
        AvlLRNIntoIter { stack: vec![self] }
    }
}

impl<T> Iterator for AvlNLRIntoIter<T> {
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        while let Some(tree) = self.stack.pop() {
            match tree {
                Null => continue,
                Tree(node) => {
                    match node.left {
                        Null => (),
                        Tree(node) => self.stack.push(Tree(node)),
                    }
                    match node.right {
                        Null => (),
                        Tree(node) => self.stack.push(Tree(node)),
                    }
                    return Some(node.val);
                }
            }
        }
        None
    }
}

impl<T> Iterator for AvlLRNIntoIter<T> {
    type Item = T;
    fn next(&mut self) -> Option<Self::Item> {
        while let Some(tree) = self.stack.pop() {
            match tree {
                Null => continue,
                Tree(node) => {
                    match node.right {
                        Null => (),
                        Tree(node) => self.stack.push(Tree(node)),
                    }
                    match node.left {
                        Null => (),
                        Tree(node) => self.stack.push(Tree(node)),
                    }
                    return Some(node.val);
                }
            }
        }
        None
    }
}

LNR 要复杂一点。

struct AvlLNRIntoIter<T> {
    stack: Vec<AvlTree<T>>,
    current: AvlTree<T>,
}

impl<T> AvlTree<T>
where
    T: Ord,
{
    fn into_iter_lnr(self) -> AvlLNRIntoIter<T> {
        AvlLNRIntoIter {
            stack: Vec::new(),
            current: self,
        }
    }
}

impl<'a, T> Iterator for AvlLNRIter<'a, T>
where
    T: Ord,
{
    type Item = &'a T;
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.current {
                Null => match self.stack.pop() {
                    Some(Tree(node)) => {
                        self.current = &node.right;
                        return Some(&node.val);
                    }
                    _ => return None,
                },
                Tree(node) => {
                    self.stack.push(self.current);
                    self.current = &node.left;
                }
            }
        }
    }
}

Iter 的实现

LNR 的 Iter 实现

struct AvlLNRIter<'a, T> {
    stack: Vec<&'a AvlTree<T>>,
    current: &'a AvlTree<T>,
}

impl<T> AvlTree<T>
where
    T: Ord,
{
    fn iter_lnr(&self) -> AvlLNRIter<T> {
        AvlLNRIter {
            stack: Vec::new(),
            current: self,
        }
    }
}

impl<'a, T> Iterator for AvlLNRIter<'a, T>
where
    T: Ord,
{
    type Item = &'a T;
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            match self.current {
                Null => match self.stack.pop() {
                    Some(Tree(node)) => {
                        self.current = &node.right;
                        return Some(&node.val);
                    }
                    _ => return None,
                },
                Tree(node) => {
                    self.stack.push(self.current);
                    self.current = &node.left;
                }
            }
        }
    }
}

会在最近几天写完这篇文章)

Thoughts? Leave a comment