~nerosnm/advela

f938891641ff403ae7816006d3fcfd993639bf26 — Søren Mortensen 1 year, 8 days ago 34c6397 master
Implement height and balance calculation
2 files changed, 283 insertions(+), 4 deletions(-)

M src/node.rs
M src/tree.rs
M src/node.rs => src/node.rs +280 -2
@@ 2,13 2,26 @@

use std::cmp::Ordering;

#[derive(Debug)]
pub struct AvlNode<T>
where
    T: Ord,
{
    /// The value contained in the node.
    val: T,
    /// The left child of this node, if there is one.
    left: Option<Box<AvlNode<T>>>,
    /// The right child of this node, if there is one.
    right: Option<Box<AvlNode<T>>>,
    /// The balance factor of this node, updated after each insertion and
    /// balance.
    ///
    /// Since the tree is balanced after each insertion, this should never be
    /// outside the range `[-2, 2]`.
    balance: i8,
    /// The height of this node, defined as the length of the path to its lowest
    /// child.
    height: u32,
}

impl<T> AvlNode<T>


@@ 21,6 34,8 @@ where
            val,
            left: None,
            right: None,
            balance: 0,
            height: 0,
        }
    }



@@ 30,6 45,7 @@ where
        V: Into<AvlNode<T>>,
    {
        self.left = Some(val.into().into());
        self.update();
    }

    /// Set the right child of this node to the given node.


@@ 38,6 54,27 @@ where
        V: Into<AvlNode<T>>,
    {
        self.right = Some(val.into().into());
        self.update();
    }

    /// Get the balance factor of this node.
    pub fn balance(&self) -> i8 {
        self.balance
    }

    /// Get the height of this node.
    pub fn height(&self) -> u32 {
        self.height
    }

    /// Update the height and balance factor of this node to reflect that a
    /// change in the heights of the child nodes may have occurred.
    pub fn update(&mut self) {
        let left = self.left.as_ref().map(|l| l.height() as i64).unwrap_or(-1);
        let right = self.right.as_ref().map(|r| r.height() as i64).unwrap_or(-1);

        self.height = (left.max(right) + 1) as u32;
        self.balance = (right - left) as i8;
    }

    /// Insert a value into the subtree that has this node as its root.


@@ 57,15 94,19 @@ where
                if let Some(ref mut left) = self.left {
                    left.insert(val);
                } else {
                    self.set_left(val);
                    self.left = Some(val.into());
                }

                self.update();
            }
            Ordering::Greater => {
                if let Some(ref mut right) = self.right {
                    right.insert(val);
                } else {
                    self.set_right(val);
                    self.right = Some(val.into());
                }

                self.update();
            }
        }
    }


@@ 187,3 228,240 @@ where
        node
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Test that the height of a node is correctly updated when children are
    /// inserted below it, starting by inserting on the left side (a
    /// lesser-valued child).
    #[test]
    fn update_height_insert_lesser_first() {
        let mut root = AvlNode::new(2);
        assert_eq!(
            root.height(),
            0,
            "the height of a node with no children should be 0"
        );

        //   2
        //  /
        // 1
        root.insert(1);
        assert_eq!(
            root.height(),
            1,
            "the height of a node should be updated when a lesser leaf is \
            inserted"
        );

        //   2
        //  /\
        // 1  3
        root.insert(3);
        assert_eq!(
            root.height(),
            1,
            "the height of a node should not change when a second child is \
            inserted with equal height to the first"
        );

        //   2
        //  /\
        // 1  3
        //    \
        //     4
        root.insert(4);
        assert_eq!(
            root.height(),
            2,
            "the height of a node should be updated when one child's height \
            becomes bigger than the other"
        );
    }

    /// Test that the height of a node is correctly updated when children are
    /// inserted below it, starting by inserting on the right side (a
    /// greater-valued child).
    #[test]
    fn update_height_insert_greater_first() {
        let mut root = AvlNode::new(2);
        assert_eq!(
            root.height(),
            0,
            "the height of a node with no children should be 0"
        );

        // 2
        // \
        //  3
        root.insert(3);
        assert_eq!(
            root.height(),
            1,
            "the height of a node should be updated when a greater leaf is \
            inserted"
        );

        //   2
        //  /\
        // 1  3
        root.insert(1);
        assert_eq!(
            root.height(),
            1,
            "the height of a node should not change when a second child is \
            inserted with equal height to the first"
        );

        //     2
        //    /\
        //   1  3
        //  /
        // 0
        root.insert(0);
        assert_eq!(
            root.height(),
            2,
            "the height of a node should be updated when one child's height \
            becomes bigger than the other"
        );
    }

    /// Test that the height of a node is correctly updated when children are
    /// directly set below it.
    #[test]
    fn update_height_set() {
        let mut root = AvlNode::new(2);
        assert_eq!(
            root.height(),
            0,
            "the height of a node with no children should be 0"
        );

        root.set_left(1);
        assert_eq!(root.height(), 1, "the height of the node should be 1");

        let height_1 = AvlNodeBuilder::new(3).right(4).build();
        root.set_right(height_1);

        assert_eq!(root.height(), 2, "the height of the node should be 2");
    }

    /// Test that the balance of a node is correctly updated when children are
    /// inserted below it, starting by inserting on the left side (a
    /// lesser-valued child).
    #[test]
    fn update_balance_insert_lesser_first() {
        let mut root = AvlNode::new(2);
        assert_eq!(
            root.balance(),
            0,
            "the balance of a node with no children should be 0"
        );

        //   2
        //  /
        // 1
        root.insert(1);
        assert_eq!(
            root.balance(),
            -1,
            "the balance of a node should be updated when a lesser leaf is \
            inserted"
        );

        //   2
        //  /\
        // 1  3
        root.insert(3);
        assert_eq!(
            root.balance(),
            0,
            "the balance of a node should be updated when a second child is \
            inserted with equal height to the first"
        );

        //   2
        //  /\
        // 1  3
        //    \
        //     4
        root.insert(4);
        assert_eq!(
            root.balance(),
            1,
            "the balance of a node should be updated when one child's height \
            becomes bigger than the other"
        );
    }

    /// Test that the balance of a node is correctly updated when children are
    /// inserted below it, starting by inserting on the right side (a
    /// greater-valued child).
    #[test]
    fn update_balance_insert_greater_first() {
        let mut root = AvlNode::new(2);
        assert_eq!(
            root.balance(),
            0,
            "the balance of a node with no children should be 0"
        );

        // 2
        // \
        //  3
        root.insert(3);
        assert_eq!(
            root.balance(),
            1,
            "the balance of a node should be updated when a greater leaf is \
            inserted"
        );

        //   2
        //  /\
        // 1  3
        root.insert(1);
        assert_eq!(
            root.balance(),
            0,
            "the balance of a node should be updated when a second child is \
            inserted with equal height to the first"
        );

        //     2
        //    /\
        //   1  3
        //  /
        // 0
        root.insert(0);
        assert_eq!(
            root.balance(),
            -1,
            "the balance of a node should be updated when one child's height \
            becomes bigger than the other"
        );
    }

    /// Test that the balance of a node is correctly updated when children are
    /// directly set below it.
    #[test]
    fn update_balance_set() {
        let mut root = AvlNode::new(2);
        assert_eq!(
            root.balance(),
            0,
            "the balance of a node with no children should be 0"
        );

        root.set_left(1);
        assert_eq!(root.balance(), -1, "the balance of the node should be -1");

        let balance_1 = AvlNodeBuilder::new(3).right(4).build();
        root.set_right(balance_1);

        assert_eq!(root.balance(), 1, "the balance of the node should be 1");
    }
}

M src/tree.rs => src/tree.rs +3 -2
@@ 4,6 4,7 @@ use std::iter::FromIterator;

use crate::node::AvlNode;

#[derive(Debug)]
pub struct AvlTree<T>
where
    T: Ord,


@@ 34,8 35,8 @@ where

    /// Insert all the values from the given iterator into the tree.
    ///
    /// If the insertions caused the tree to become unbalanced, it will be
    /// balanced after all of the insertions take place.
    /// For each insertion, if it caused the tree to become unbalanced, it will
    /// be balanced afterwards.
    pub fn insert_all<I>(&mut self, iter: I)
    where
        I: IntoIterator<Item = T>,