avl树

An AVL Tree is a self-balancing binary search tree. The heights of any two sibling nodes must differ by at most one; the tree may rebalance itself after insertion or deletion to uphold this property.

Properties

  • Worst/Average time complexity for basic operations: O(log n)
  • Worst/Average space complexity: O(n)
  1. use std::{
  2.     cmp::{max, Ordering},
  3.     iter::FromIterator,
  4.     mem,
  5.     ops::Not,
  6. };
  8. /// An internal node of an `AVLTree`.
  9. struct AVLNode<T: Ord> {
  10.     value: T,
  11.     height: usize,
  12.     left: Option<Box<AVLNode<T>>>,
  13.     right: Option<Box<AVLNode<T>>>,
  14. }
  16. /// A set based on an AVL Tree.
  17. ///
  18. /// An AVL Tree is a self-balancing binary search tree. It tracks the height of each node
  19. /// and performs internal rotations to maintain a height difference of at most 1 between
  20. /// each sibling pair.
  21. pub struct AVLTree<T: Ord> {
  22.     root: Option<Box<AVLNode<T>>>,
  23.     length: usize,
  24. }
  26. /// Refers to the left or right subtree of an `AVLNode`.
  27. #[derive(Clone, Copy)]
  28. enum Side {
  29.     Left,
  30.     Right,
  31. }
  33. impl<T: Ord> AVLTree<T> {
  34.     /// Creates an empty `AVLTree`.
  35.     pub fn new() -> AVLTree<T> {
  36.         AVLTree {
  37.             root: None,
  38.             length: 0,
  39.         }
  40.     }
  42.     /// Returns `true` if the tree contains a value.
  43.     pub fn contains(&self, value: &T) -> bool {
  44.         let mut current = &self.root;
  45.         while let Some(node) = current {
  46.             current = match value.cmp(&node.value) {
  47.                 Ordering::Equal => return true,
  48.                 Ordering::Less => &node.left,
  49.                 Ordering::Greater => &node.right,
  50.             }
  51.         }
  52.         false
  53.     }
  55.     /// Adds a value to the tree.
  56.     ///
  57.     /// Returns `true` if the tree did not yet contain the value.
  58.     pub fn insert(&mut self, value: T) -> bool {
  59.         let inserted = insert(&mut self.root, value);
  60.         if inserted {
  61.             self.length += 1;
  62.         }
  63.         inserted
  64.     }
  66.     /// Removes a value from the tree.
  67.     ///
  68.     /// Returns `true` if the tree contained the value.
  69.     pub fn remove(&mut self, value: &T) -> bool {
  70.         let removed = remove(&mut self.root, value);
  71.         if removed {
  72.             self.length -= 1;
  73.         }
  74.         removed
  75.     }
  77.     /// Returns the number of values in the tree.
  78.     pub fn len(&self) -> usize {
  79.         self.length
  80.     }
  82.     /// Returns `true` if the tree contains no values.
  83.     pub fn is_empty(&self) -> bool {
  84.         self.length == 0
  85.     }
  87.     /// Returns an iterator that visits the nodes in the tree in order.
  88.     fn node_iter(&self) -> NodeIter<T> {
  89.         let cap = self.root.as_ref().map_or(0, |n| n.height);
  90.         let mut node_iter = NodeIter {
  91.             stack: Vec::with_capacity(cap),
  92.         };
  93.         // Initialize stack with path to leftmost child
  94.         let mut child = &self.root;
  95.         while let Some(node) = child {
  96.             node_iter.stack.push(node.as_ref());
  97.             child = &node.left;
  98.         }
  99.         node_iter
  100.     }
  102.     /// Returns an iterator that visits the values in the tree in ascending order.
  103.     pub fn iter(&self) -> Iter<T> {
  104.         Iter {
  105.             node_iter: self.node_iter(),
  106.         }
  107.     }
  108. }
  110. /// Recursive helper function for `AVLTree` insertion.
  111. fn insert<T: Ord>(tree: &mut Option<Box<AVLNode<T>>>, value: T) -> bool {
  112.     if let Some(node) = tree {
  113.         let inserted = match value.cmp(&node.value) {
  114.             Ordering::Equal => false,
  115.             Ordering::Less => insert(&mut node.left, value),
  116.             Ordering::Greater => insert(&mut node.right, value),
  117.         };
  118.         if inserted {
  119.             node.rebalance();
  120.         }
  121.         inserted
  122.     } else {
  123.         *tree = Some(Box::new(AVLNode {
  124.             value,
  125.             height: 1,
  126.             left: None,
  127.             right: None,
  128.         }));
  129.         true
  130.     }
  131. }
  133. /// Recursive helper function for `AVLTree` deletion.
  134. fn remove<T: Ord>(tree: &mut Option<Box<AVLNode<T>>>, value: &T) -> bool {
  135.     if let Some(node) = tree {
  136.         let removed = match value.cmp(&node.value) {
  137.             Ordering::Less => remove(&mut node.left, value),
  138.             Ordering::Greater => remove(&mut node.right, value),
  139.             Ordering::Equal => {
  140.                 *tree = match (node.left.take(), node.right.take()) {
  141.                     (None, None) => None,
  142.                     (Some(b), None) | (None, Some(b)) => Some(b),
  143.                     (Some(left), Some(right)) => Some(merge(left, right)),
  144.                 };
  145.                 return true;
  146.             }
  147.         };
  148.         if removed {
  149.             node.rebalance();
  150.         }
  151.         removed
  152.     } else {
  153.         false
  154.     }
  155. }
  157. /// Merges two trees and returns the root of the merged tree.
  158. fn merge<T: Ord>(left: Box<AVLNode<T>>, right: Box<AVLNode<T>>) -> Box<AVLNode<T>> {
  159.     let mut op_right = Some(right);
  160.     // Guaranteed not to panic since right has at least one node
  161.     let mut root = take_min(&mut op_right).unwrap();
  162.     root.left = Some(left);
  163.     root.right = op_right;
  164.     root.rebalance();
  165.     root
  166. }
  168. /// Removes the smallest node from the tree, if one exists.
  169. fn take_min<T: Ord>(tree: &mut Option<Box<AVLNode<T>>>) -> Option<Box<AVLNode<T>>> {
  170.     if let Some(mut node) = tree.take() {
  171.         // Recurse along the left side
  172.         if let Some(small) = take_min(&mut node.left) {
  173.             // Took the smallest from below; update this node and put it back in the tree
  174.             node.rebalance();
  175.             *tree = Some(node);
  176.             Some(small)
  177.         } else {
  178.             // Take this node and replace it with its right child
  179.             *tree = node.right.take();
  180.             Some(node)
  181.         }
  182.     } else {
  183.         None
  184.     }
  185. }
  187. impl<T: Ord> AVLNode<T> {
  188.     /// Returns a reference to the left or right child.
  189.     fn child(&self, side: Side) -> &Option<Box<AVLNode<T>>> {
  190.         match side {
  191.             Side::Left => &self.left,
  192.             Side::Right => &self.right,
  193.         }
  194.     }
  196.     /// Returns a mutable reference to the left or right child.
  197.     fn child_mut(&mut self, side: Side) -> &mut Option<Box<AVLNode<T>>> {
  198.         match side {
  199.             Side::Left => &mut self.left,
  200.             Side::Right => &mut self.right,
  201.         }
  202.     }
  204.     /// Returns the height of the left or right subtree.
  205.     fn height(&self, side: Side) -> usize {
  206.         self.child(side).as_ref().map_or(0, |n| n.height)
  207.     }
  209.     /// Returns the height difference between the left and right subtrees.
  210.     fn balance_factor(&self) -> i8 {
  211.         let (left, right) = (self.height(Side::Left), self.height(Side::Right));
  212.         if left < right {
  213.             (right - left) as i8
  214.         } else {
  215.             -((left - right) as i8)
  216.         }
  217.     }
  219.     /// Recomputes the `height` field.
  220.     fn update_height(&mut self) {
  221.         self.height = 1 + max(self.height(Side::Left), self.height(Side::Right));
  222.     }
  224.     /// Performs a left or right rotation.
  225.     fn rotate(&mut self, side: Side) {
  226.         let mut subtree = self.child_mut(!side).take().unwrap();
  227.         *self.child_mut(!side) = subtree.child_mut(side).take();
  228.         self.update_height();
  229.         // Swap root and child nodes in memory
  230.         mem::swap(self, subtree.as_mut());
  231.         // Set old root (subtree) as child of new root (self)
  232.         *self.child_mut(side) = Some(subtree);
  233.         self.update_height();
  234.     }
  236.     /// Performs left or right tree rotations to balance this node.
  237.     fn rebalance(&mut self) {
  238.         self.update_height();
  239.         let side = match self.balance_factor() {
  240.             -2 => Side::Left,
  241.             2 => Side::Right,
  242.             _ => return,
  243.         };
  244.         let subtree = self.child_mut(side).as_mut().unwrap();
  245.         // Left-Right and Right-Left require rotation of heavy subtree
  246.         if let (Side::Left, 1) | (Side::Right, -1) = (side, subtree.balance_factor()) {
  247.             subtree.rotate(side);
  248.         }
  249.         // Rotate in opposite direction of heavy side
  250.         self.rotate(!side);
  251.     }
  252. }
  254. impl<T: Ord> Default for AVLTree<T> {
  255.     fn default() -> Self {
  256.         Self::new()
  257.     }
  258. }
  260. impl Not for Side {
  261.     type Output = Side;
  263.     fn not(self) -> Self::Output {
  264.         match self {
  265.             Side::Left => Side::Right,
  266.             Side::Right => Side::Left,
  267.         }
  268.     }
  269. }
  271. impl<T: Ord> FromIterator<T> for AVLTree<T> {
  272.     fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
  273.         let mut tree = AVLTree::new();
  274.         for value in iter {
  275.             tree.insert(value);
  276.         }
  277.         tree
  278.     }
  279. }
  281. /// An iterator over the nodes of an `AVLTree`.
  282. ///
  283. /// This struct is created by the `node_iter` method of `AVLTree`.
  284. struct NodeIter<'a, T: Ord> {
  285.     stack: Vec<&'a AVLNode<T>>,
  286. }
  288. impl<'a, T: Ord> Iterator for NodeIter<'a, T> {
  289.     type Item = &'a AVLNode<T>;
  291.     fn next(&mut self) -> Option<Self::Item> {
  292.         if let Some(node) = self.stack.pop() {
  293.             // Push left path of right subtree to stack
  294.             let mut child = &node.right;
  295.             while let Some(subtree) = child {
  296.                 self.stack.push(subtree.as_ref());
  297.                 child = &subtree.left;
  298.             }
  299.             Some(node)
  300.         } else {
  301.             None
  302.         }
  303.     }
  304. }
  306. /// An iterator over the items of an `AVLTree`.
  307. ///
  308. /// This struct is created by the `iter` method of `AVLTree`.
  309. pub struct Iter<'a, T: Ord> {
  310.     node_iter: NodeIter<'a, T>,
  311. }
  313. impl<'a, T: Ord> Iterator for Iter<'a, T> {
  314.     type Item = &'a T;
  316.     fn next(&mut self) -> Option<&'a T> {
  317.         match self.node_iter.next() {
  318.             Some(node) => Some(&node.value),
  319.             None => None,
  320.         }
  321.     }
  322. }
  324. #[cfg(test)]
  325. mod tests {
  326.     use super::AVLTree;
  328.     /// Returns `true` if all nodes in the tree are balanced.
  329.     fn is_balanced<T: Ord>(tree: &AVLTree<T>) -> bool {
  330.         tree.node_iter()
  331.             .all(|n| (-1..=1).contains(&n.balance_factor()))
  332.     }
  334.     #[test]
  335.     fn len() {
  336.         let tree: AVLTree<_> = (1..4).collect();
  337.         assert_eq!(tree.len(), 3);
  338.     }
  340.     #[test]
  341.     fn contains() {
  342.         let tree: AVLTree<_> = (1..4).collect();
  343.         assert!(tree.contains(&1));
  344.         assert!(!tree.contains(&4));
  345.     }
  347.     #[test]
  348.     fn insert() {
  349.         let mut tree = AVLTree::new();
  350.         // First insert succeeds
  351.         assert!(tree.insert(1));
  352.         // Second insert fails
  353.         assert!(!tree.insert(1));
  354.     }
  356.     #[test]
  357.     fn remove() {
  358.         let mut tree: AVLTree<_> = (1..8).collect();
  359.         // First remove succeeds
  360.         assert!(tree.remove(&4));
  361.         // Second remove fails
  362.         assert!(!tree.remove(&4));
  363.     }
  365.     #[test]
  366.     fn sorted() {
  367.         let tree: AVLTree<_> = (1..8).rev().collect();
  368.         assert!((1..8).eq(tree.iter().map(|&x| x)));
  369.     }
  371.     #[test]
  372.     fn balanced() {
  373.         let mut tree: AVLTree<_> = (1..8).collect();
  374.         assert!(is_balanced(&tree));
  375.         for x in 1..8 {
  376.             tree.remove(&x);
  377.             assert!(is_balanced(&tree));
  378.         }
  379.     }
  380. }