How to write a delete function for a binary tree in Rust?

user_5784021 0
Viewed 776

I'm trying to learn some Rust and started with implementing a binary tree. While doing insertions and in-order traversal seem easy enough I'm having some trouble with the borrow-checker in the delete method.

The code so far is:

use std::cmp::{Ordering, PartialOrd};
use std::fmt::Display;

struct Node<T> {
    value: T,
    left: Option<Box<Node<T>>>,
    right: Option<Box<Node<T>>>,
}

impl<T: Display + PartialOrd> Node<T> {
    fn new(val: T) -> Self {
        Node {
            value: val,
            left: None,
            right: None,
        }
    }

    fn print_inorder(&self) {
        print!("[");
        self.inorder(|x| print!("{}, ", x));
        print!("]\n")
    }

    fn inorder(&self, f: fn(&T)) {
        if let Some(ref x) = self.left {
            (*x).inorder(f);
        }
        f(&self.value);
        if let Some(ref x) = self.right {
            (*x).inorder(f);
        }
    }

    fn insert(&mut self, val: T) {
        let mut node = self;

        loop {
            if val == node.value {
                return;
            }

            let child = match val.partial_cmp(&node.value).expect("Key comparison failed") {
                Ordering::Less => &mut node.left,
                Ordering::Equal => return,
                Ordering::Greater => &mut node.right,
            };

            match child {
                Some(ref mut c) => node = c,
                None => { *child = Some(Box::new(Node::new(val))); return}
            }
        }
    }

    fn delete(&mut self, val: T) -> bool {
        if self.value == val {
            unimplemented!{"Cannot remove root (yet) :/"};
        }

        let mut node = self;

        loop {
            if val < node.value {
                if let Some(ref mut c) = node.left {

                    if c.value == val {
                        if c.left.is_none() && c.right.is_none() {
                            // Error: cannot assign to node.left here
                            node.left = None;
                        }
                        else if c.left.is_some() && c.right.is_none() {
                            // Error: again cannot assign to node.left but cannot even access c.left
                            node.left = c.left;
                        }
                        else if c.left.is_none() && c.right.is_some() {
                            node.left = c.right;
                        }
                        else {
                            // TODO: walk through child tree and get rightmost element
                        }
                        return true;
                    } else {
                        node = c;
                    }
                }
                else {
                    return false;
                }
            }
        }
    }
}

struct Tree<T> {
    root: Option<Node<T>>
}

impl<T: PartialOrd + Display> Tree<T> {
    fn new() -> Self {
        Tree { root: None }
    }

    fn insert(&mut self, val: T) {
        match self.root {
            Some(ref mut n) => n.insert(val),
            None => self.root = Some(Node::new(val))
        }
    }

    fn print(&self) {
        match self.root {
            Some(ref n) => n.print_inorder(),
            None => println!("[]")
        }
    }
}

fn main() {
    println!("Hello, world!");

    let mut t = Tree::<i64>::new();
    t.print();
    t.insert(14);
    t.print();
    t.insert(7);
    t.insert(21);
    t.print();
}

So in the delete method I cannot assign to the nodes I traversed downwards. But to delete a Node I have to assign a different value to its parent.

I found this implementation but it doesn't compile for me (again with borrow-checker problems).

rust
1 Answers

We can use a recursive solution to dodge many of the borrowing issues that the iterative solution faces, as a starting point.

In order to be able to delete the root node, the key idea is to ask the delete function to return the new sub-tree rooted in the node; a slight modification is necessary in Tree, to harmonize things:

struct Tree<T> {
    root: Option<Box<Node<T>>>
}

I added the Box layer, because Node contains Option<Box<Node<T>>>; this allows us handling root like another child node.

Then, I code the recursive delete:

fn delete(mut this: Box<Node<T>>, target: &T) -> Option<Box<Node<T>>> {
    if target < &this.value {
        if let Some(left) = this.left.take() {
            this.left = Self::delete(left, target);
        }
        return Some(this);
    }

    if target > &this.value {
        if let Some(right) = this.right.take() {
            this.right = Self::delete(right, target);
        }
        return Some(this);
    }

    assert!(target == &this.value, "Faulty Ord implementation for T");

    match (this.left.take(), this.right.take()) {
        (None, None) => None,
        (Some(left), None) => Some(left),
        (None, Some(right)) => Some(right),
        (Some(mut left), Some(right)) => {
            if let Some(mut rightmost) = left.rightmost_child() {
                rightmost.left = Some(left);
                rightmost.right = Some(right);
                Some(rightmost)
            } else {
                left.right = Some(right);
                Some(left)
            }
        },
    }
}

//  Returns the rightmost child, unless the node itself is that child.
fn rightmost_child(&mut self) -> Option<Box<Node<T>>> {
    match self.right {
        Some(ref mut right) =>  {
            if let Some(t) = right.rightmost_child() {
                Some(t)
            } else {
                let mut r = self.right.take();
                if let Some(ref mut r) = r {
                    self.right = std::mem::replace(&mut r.left, None);
                }
                r
            }
        },
        None => None
    }
}

Which is called from Tree via:

fn delete(&mut self, target: &T) {
    if let Some(root) = self.root.take() {
        self.root = Node::delete(root, target);
    }
}

The complete code is available on the playground.

user_1471920
Related