最近在折腾一些基础知识,想趁机写一些博客。希望能坚持下来吧。
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; } } } } }
会在最近几天写完这篇文章)