Morris traversal in Rust

morris_traversal.rs

use std::iter::FromIterator;
use std::cell::RefCell;
use std::rc::Rc;
use std::collections::VecDeque;

fn main() {
    print!("Morris中序遍历：");
    let nums = vec![6, 2, 7, 1, 4, -1, 9, -1, -1, 3, 5, 8, -1];
    let tree = Tree::from_iter(nums);
    let nums = tree.inorder_morris_traversal();
    println!("{:?}", nums);

    print!("Morris先序遍历：");
    let nums = vec![1, 2, 7, 3, 4, -1, 8, -1, -1, 5, 6, 9, -1];
    let tree = Tree::from_iter(nums);
    let nums = tree.preorder_morris_traversal();
    println!("{:?}", nums);

    print!("Morris后序遍历：");
    let nums = vec![9, 5, 8, 1, 4, -1, 7, -1, -1, 2, 3, 6, -1];
    let tree = Tree::from_iter(nums);
    let nums = tree.postorder_morris_traversal();
    println!("{:?}", nums);
}

pub struct Tree(Option<Rc<RefCell<TreeNode>>>);

impl FromIterator<i32> for Tree {
    fn from_iter<T: IntoIterator<Item = i32>>(iter: T) -> Self {
        let mut iter = iter.into_iter();
        let head = match iter.next() {
            Some(val) => TreeNode::new(val),
            None => return Tree(None),
        };
        let mut queue = VecDeque::new();
        queue.push_back(head.clone());
        loop {
            let node = queue.pop_front().unwrap();
            let mut node = node.borrow_mut();
            match iter.next() {
                Some(-1) => {}
                Some(val) => {
                    let left = TreeNode::new(val);
                    queue.push_back(left.clone());
                    node.left = Some(left);
                }
                None => break,
            }
            match iter.next() {
                Some(-1) => {}
                Some(val) => {
                    let right = TreeNode::new(val);
                    queue.push_back(right.clone());
                    node.right = Some(right);
                }
                None => break,
            }
        }
        Tree(Some(head))
    }
}

impl IntoIterator for Tree {
    type Item = i32;
    type IntoIter = TreeIter;

    fn into_iter(self) -> Self::IntoIter {
        let mut queue = VecDeque::new();
        match self.0 {
            Some(root) => queue.push_back(root),
            None => {}
        }
        TreeIter { queue }
    }
}

impl Tree {
    pub fn inorder_morris_traversal(&self) -> Vec<i32> {
        let mut nums = Vec::new();
        let mut curr = self.0.clone();
        while let Some(curr_node) = curr {
            let ref_curr_node = curr_node.borrow();
            let mut prev = match &ref_curr_node.left {
                Some(left) => left.clone(),
                None => {
                    // 如果当前节点的左孩子为空，则**输出当前节点**并将其右孩子作为当前节点。
                    nums.push(ref_curr_node.val);
                    curr = ref_curr_node.right.clone();
                    continue;
                }
            };
            // 如果当前节点的左孩子不为空，在当前节点的左子树中找到当前节点在中序遍历下的前驱节点。
            loop {
                let ref_prev = prev.borrow();
                let right = match &ref_prev.right {
                    None => break,
                    Some(right) if Rc::ptr_eq(right, &curr_node) => break,
                    Some(right) => right.clone(),
                };
                drop(ref_prev);
                prev = right;
            }
            let mut prev = prev.borrow_mut();
            match prev.right {
                None => {
                    // 如果前驱节点的右孩子为空，将它的右孩子设置为当前节点。当前节点更新为当前节点的左孩子。
                    prev.right = Some(curr_node.clone());
                    curr = ref_curr_node.left.clone();
                }
                Some(_) => {
                    // 如果前驱节点的右孩子为当前节点，将它的右孩子重新设为空（恢复树的形状）。
                    // **输出当前节点**。当前节点更新为当前节点的右孩子。
                    prev.right = None;
                    nums.push(ref_curr_node.val);
                    curr = ref_curr_node.right.clone();
                }
            }
        }
        nums
    }

    pub fn preorder_morris_traversal(&self) -> Vec<i32> {
        let mut nums = Vec::new();
        let mut curr = self.0.clone();
        while let Some(curr_node) = curr {
            let ref_curr_node = curr_node.borrow();
            let mut prev = match &ref_curr_node.left {
                Some(left) => left.clone(),
                None => {
                    // 如果当前节点的左孩子为空，则**输出当前节点**并将其右孩子作为当前节点。
                    nums.push(ref_curr_node.val);
                    curr = ref_curr_node.right.clone();
                    continue;
                }
            };
            // 如果当前节点的左孩子不为空，在当前节点的左子树中找到当前节点在中序遍历下的前驱节点。
            loop {
                let ref_prev = prev.borrow();
                let right = match &ref_prev.right {
                    None => break,
                    Some(right) if Rc::ptr_eq(right, &curr_node) => break,
                    Some(right) => right.clone(),
                };
                drop(ref_prev);
                prev = right;
            }
            let mut prev = prev.borrow_mut();
            if prev.right.is_none() {
                // 如果前驱节点的右孩子为空，将它的右孩子设置为当前节点。
                // 当前节点更新为当前节点的左孩子。**输出当前节点**。
                nums.push(ref_curr_node.val);
                prev.right = Some(curr_node.clone());
                curr = ref_curr_node.left.clone();
            } else {
                // 如果前驱节点的右孩子为当前节点，将它的右孩子重新设为空（恢复树的形状）。
                // 当前节点更新为当前节点的右孩子。
                prev.right = None;
                curr = ref_curr_node.right.clone();
            }
        }
        nums
    }

    pub fn postorder_morris_traversal(&self) -> Vec<i32> {
        fn reverse(from: &Rc<RefCell<TreeNode>>, to: &Rc<RefCell<TreeNode>>) {
            if Rc::ptr_eq(from, to) {
                return;
            }
            let mut x = from.clone();
            let mut y = from.borrow().right.as_ref().unwrap().clone();
            loop {
                let z = y.borrow().right.as_ref().unwrap().clone();
                y.borrow_mut().right = Some(x);
                x = y;
                y = z;
                if Rc::ptr_eq(&x, to) {
                    break;
                }
            }
        }
        fn push_reverse(
            from: &Rc<RefCell<TreeNode>>,
            to: &Rc<RefCell<TreeNode>>,
            nums: &mut Vec<i32>,
        ) {
            reverse(from, to);
            let mut curr = to.clone();
            loop {
                nums.push(curr.borrow().val);
                if Rc::ptr_eq(&curr, from) {
                    break;
                }
                let next = curr.borrow().right.as_ref().unwrap().clone();
                curr = next;
            }
            reverse(to, from);
        }

        let mut nums = Vec::new();
        // 设置当前节点为临时结点
        let dump = TreeNode::new(0);
        dump.borrow_mut().left = self.0.clone();
        let mut curr = Some(dump);
        while let Some(curr_node) = curr {
            let ref_curr_node = curr_node.borrow();
            let mut prev = match &ref_curr_node.left {
                Some(left) => left.clone(),
                None => {
                    // 如果当前节点的左孩子为空，则将其右孩子作为当前节点。
                    curr = ref_curr_node.right.clone();
                    continue;
                }
            };
            // 如果当前节点的左孩子不为空，在当前节点的左子树中找到当前节点在中序遍历下的前驱节点。
            loop {
                let ref_prev = prev.borrow();
                let right = match &ref_prev.right {
                    None => break,
                    Some(right) if Rc::ptr_eq(right, &curr_node) => break,
                    Some(right) => right.clone(),
                };
                drop(ref_prev);
                prev = right;
            }
            let mut ref_prev = prev.borrow_mut();
            if ref_prev.right.is_none() {
                // 如果前驱节点的右孩子为空，将它的右孩子设置为当前节点。当前节点更新为当前节点的左孩子。
                ref_prev.right = Some(curr_node.clone());
                curr = ref_curr_node.left.clone();
            } else {
                // 如果前驱节点的右孩子为当前节点，将它的右孩子重新设为空（恢复树的形状）。
                // **倒序输出当前节点左孩子到前序结点的所有节点**当前节点更新为当前节点的右孩子。
                drop(ref_prev);
                push_reverse(ref_curr_node.left.as_ref().unwrap(), &prev, &mut nums);
                prev.borrow_mut().right = None;
                curr = ref_curr_node.right.clone();
            }
        }
        nums
    }
}

#[derive(Debug)]
pub struct TreeNode {
    pub val: i32,
    pub left: Option<Rc<RefCell<TreeNode>>>,
    pub right: Option<Rc<RefCell<TreeNode>>>,
}

impl TreeNode {
    pub fn new(val: i32) -> Rc<RefCell<Self>> {
        Rc::new(RefCell::new(TreeNode {
            val,
            left: None,
            right: None,
        }))
    }
}

pub struct TreeIter {
    queue: VecDeque<Rc<RefCell<TreeNode>>>,
}

impl Iterator for TreeIter {
    type Item = i32;

    fn next(&mut self) -> Option<Self::Item> {
        match self.queue.pop_front() {
            Some(node) => {
                let node = node.borrow();
                let val = node.val;
                if let Some(left) = node.left.as_ref() {
                    self.queue.push_back(left.clone());
                }
                if let Some(right) = node.right.as_ref() {
                    self.queue.push_back(right.clone());
                }
                Some(val)
            }
            None => None,
        }
    }
}