Rust 二叉树实现与遍历

二叉树(Binary Tree) 是非常基础的数据结构。平衡树可以让树的查找,更新,插入,删除都是O(logN)的复杂度。

二叉树的基本实现是比较简单的。然而对于Rust,因为所有权(ownership)和借用检查(borrow check) 机制,会比普通的语言稍微"麻烦"一点。leetcode 上有很多二叉树的题目,本文采用leetcode定义的 tree 节点类型。一方面是 leetcode 的定义可以方便刷题,另一方面会涉及到rust语言特性,可以练习。

Note: leetcode上定义的rust treenode节点,并不是特别match rust的特性,原因参考

本文的主要内容如下:

  • 二叉树节点定义,使用rust 如何定义二叉树节点
  • 二叉树的遍历,dfs(先序,中序,后序),bfs(层序) 的写法,递归和迭代的异同
  • 二叉树的树形打印,方便直观的看到一颗二叉树的拓扑形状
  • 根据一个数组构建一颗二叉树,即Leetcode二叉树的输入,方便离线环境刷题

树节点定义

Leetcode 几乎给出了其支持语言的树节点的定义。如 python 的节点

class TreeNode:
    def __init__(self, x):
        self.val = x
        self.left = None
        self.right = None

python节点很简单,定义一个 class,属性分别是节点值,左右子树。

golang 的节点

type TreeNode struct {
    Val int
    Left *TreeNode
    Right *TreeNode
}

golang节点也不复杂,定义一个结构体,也是三个字段,分表是数据值,左右子树的节点指针。之所以左右子树用指针。TreeNode是递归定义,TreeNode如果还没声明,是无法在其字段中定义声明的。

然而 rust的节点定义,却看起来稍微复杂:

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

impl TreeNode {
  #[inline]
  pub fn new(val: i32) -> Self {
    TreeNode {
      val,
      left: None,
      right: None
    }
  }
}

对于 #[derive(Debug, PartialEq, Eq)] 是rust的一种注解,表示默认实现的trait。这里表面这个结构体可以打印(Debug)到控制台以及可以比较大小(PartialEq, Eq)。

然后定义 TreeNode 节点的时候,有了三层。

  1. Option 一种枚举,有两个变体 Some(T) 和 None。有值的时候使用 Some 包装,无值的时候使用 None。类似 python 的 None 和 golang 的nil,但与它们有本质的区别。
  2. Rc 是一种智能指针,可以共享所有权。rust 的变量和其值是绑定关系,每个值都有管理其所有权的变量。Rc是 (Reference Count)引用计数,它可以有多个所有权引用。但是Rc包装的内容不可变。
  3. 因为Rc包装的内容不能改变,而 treenode 会变更其左右子树,因此 RefCell 提供了内部可变性。

Option比较好理解,下面分别对剩下两个原因进行代码演示。为了方便演示,定义下面的简化结构。

#[derive(Debug)]
struct TreeNode {
    val: i32
}
    let t1 = TreeNode{val:1};
    println!("{:?}", t1); // 输出: TreeNode { val: 1 }

定义一个 t1 变量,然后绑定一个 TreeNode 实例。通常 tree 的遍历需要使用栈。

    let t1 = TreeNode{val:1};
    println!("{:?}", t1);

    let mut stack = vec![];
    stack.push(t1);

    println!("{:?}", t1);

让 t1 进栈,然后再打印 t1。很不幸,rust编译器将拒绝。原因就是 stack.push(t1),将所有权转移到了push函数内。t1 变量就等同于"失效"了, println!再使用 t1 就会报错。下面是 rust 编译错误

17 |     let t1 = TreeNode{val:1};
   |         -- move occurs because `t1` has type `TreeNode`, which does not implement the `Copy` trait
...
21 |     stack.push(t1);
   |                -- value moved here
22 | 
23 |     println!("{:?}", t1);
   |                      ^^ value borrowed here after move

使用 Rc 就能保持多个所有权引用。

    let t1 = Rc::new(TreeNode{val:1});  // 使用 Rc 包装 TreeNode

    let mut stack = vec![];
    stack.push(t1.clone());             // 增加 t1 的引用计数,其堆上的数据不变。

    println!("{:?}", t1); // TreeNode { val: 1 }
    println!("{:?}", Rc::strong_count(&t1))  // 2  

通过 Rc 可以看到完美的解决了多所有权引用问题。t1.clone() 还可以写成 Rc::clone(&t1)。后者其实更符合rust习惯,因为可以和其他数据结构的 .clone 方法区分。通常其他结构的 .clone 方法是深复制。而 Rc 的 .clone 是浅复制,只增加 Rc 指针的引用计数,不会复制堆上的数据。

到目前为止,一切顺利。实际上,我们对树节点除了一般的查找,还会更新。也就是会更新 TreeNode内部字段的值。下面的代码将不会编译通过:


    let mut t1 = Rc::new(TreeNode{val:1});

    let mut stack = vec![];
    stack.push(t1.clone());

    println!("{:?}", t1); // TreeNode { val: 1 }
    println!("{:?}", Rc::strong_count(&t1));  // 2

    t1.val = 12;
    // t1.deref().val = 12;
    // t1.borrow() = 12;
    // t1.borrow_mut() = 12;


t1.val 可以读到值。因为 Rc 实现了 Deref trai,会自动解引用。但是不代表可以 t1.val 直接赋值修改。变量 mut t1 的修饰表示 t1 可以重新赋值别的值,但是还是不能修改其自生内部字段的值。同样的,t1.deref 和 t1.borrow t1.borrow_mut 得到的 &TreeNode ,但是依然都无法修改 val 的值。

想要修改 TreeNode 内部字段的值,就需要 RefCell 结构。

    let t1 = Rc::new(RefCell::new(TreeNode{val:1}));   // 增加 RefCell 包装

    let mut stack = vec![];
    stack.push(t1.clone());

    println!("{:?}", t1); // TreeNode { val: 1 }
    println!("{:?}", Rc::strong_count(&t1));  // 2

    t1.borrow_mut().val = 12;                 // 通过 borrow_mut 方法获取 RefMut 类型,后者可以直接修改 val 的值。

    println!("{:?}", t1.borrow().val);  // 

RefCell 结构通过 borrow/borrow_mut 方法可以得到 Ref/RefMut 类型。后者就等同于 Treenode,可以直接改变其内部字段。并且是 t1 的内部可变,其本身还是 不可变的。

综上所述,Rc 用于多所有权的引用,RefCell 用于修改 struct 内部字段的值。所以 Leetcode 如此定义 TreeNode。

树的遍历

二叉树的遍历主要有两大类,深度优先搜索(DFS)和广度优先(BFS)。前者是指沿着树的路径先遍历到叶子(Leaf)节点,然后再遍历其他路径。后者是指先遍历每一层的节点,然后层层递进。

二叉树有两个很有用的属性,节点的高度和深度。高度就是节点到最深叶子节点的距离。其深度是节点到树根的距离(Leetcode 的定义,实际上路径数和节点数两个量都可以)。

深度优先

如同树节点的定义一样,二叉树是节点的集合,一个节点也是二叉树。二叉树可以看成:左子树,根,右子树的最基本单元。


traversal.png

根据 这三者的遍历顺序可以分为三者遍历方法。

  • 先序(前序)遍历(preorder traversal): 根 -> 左子树 -> 右子树
  • 中序遍历(inorder traversal):左子树 -> 根 -> 右子树
  • 后序遍历(inorder traversal):左子树 -> 右子树 -> 根

递归遍历

三者遍历的递归写法很简单。

先序遍历

preorder.png

例如下面的 Python3代码。

def preorder(root: TreeNode):
    # 递归基,节点不存在,直接返回
    if root is None:
        return 
    # 先访问根节点的值
    yield root.val 
    # 递归遍历左子树,即以左子树为树根进行二叉树遍历
    yield from preorder(root.left)
    # 递归遍历左子树,即以左子树为树根进行二叉树遍历
    yield from preorder(root.left)

之所以这里使用python,是因为py的代码和算法伪代码很相似,后面再介绍 rust的实现。

中序遍历

inorder.png
def inorder(root: TreeNode):
    # 递归基,节点不存在,直接返回
    if root is None:
        return 

    # 先递归遍历左子树,即以左子树为树根进行二叉树遍历
    yield from inorder(root.left)
    # 访问根节点的值
    yield root.val 
    # 递归遍历左子树,即以左子树为树根进行二叉树遍历
    yield from inorder(root.left)

后序遍历

postorder.png
def postorder(root: TreeNode):
    # 递归基,节点不存在,直接返回
    if root is None:
        return 

    # 先递归遍历左子树,即以左子树为树根进行二叉树遍历
    yield from postorder(root.left)

    # 递归遍历左子树,即以左子树为树根进行二叉树遍历
    yield from postorder(root.left)

    # 最后访问根节点的值
    yield root.val 

从上面的代码可以看出,二叉树的DFS遍历的差别只是再访问根的节点的顺序,这本身也是其定义。先根,中根,后根。因此可以统一下面的伪代码


def tree_order(node):
    if node is None:
        return 

    # preorder vist node
    tree_order(node.left)
    # inorder vist node
    tree_order(node.right)
    # postsorder vist node

rust 实现

leetcode 144. 二叉树的前序遍历 是二叉树的先序遍历。套用上面介绍的算法,rust代码如下

use std::rc::Rc;
use std::cell::RefCell;

impl Solution {
    pub fn preorder_traversal(root: Option>>) -> Vec {
        fn preorder_dfs(node: &Option>>, ans: &mut Vec) {
            // 递归基
            if node.is_none(){
                return
            }
            // 访问根节点
            ans.push(node.as_ref().unwrap().borrow().val);
            // 访问左子树
            preorder_dfs(&node.as_ref().unwrap().borrow().left, ans);
            // 访问右子树
            preorder_dfs(&node.as_ref().unwrap().borrow().right, ans);
        }
        let mut ans = vec![];
        preorder_dfs(&root, &mut ans);
        ans
    }
}

因为 node 是 &Option>> 类型,因此需要先提取 option内的 rc,然后通过 borrow 活动 Ref类型的 TreeNode。上面代码写法和算法模板一致,但不是 rust 常用的写法。下面中序遍历用 rust 模式匹配进行重构

94. 二叉树的中序遍历。

impl Solution {
    pub fn inorder_traversal(root: Option>>) -> Vec {

        fn inorder_dfs(node: &Option>>, ans: &mut Vec){
            match node {
                None => {},
                Some(x) => {
                    inorder_dfs(&x.borrow().left, ans);
                    ans.push(x.borrow().val);
                    inorder_dfs(&x.borrow().right, ans);
                },

            }
        }

        let mut ans = vec![];
        inorder_dfs(&root, &mut ans);
        ans
    }
}

上面的代码使用了 match 模式匹配。其写法依然和算法模板很像。node是 &Option, 通过 match模式匹配,None 相当于 递归基,Some 变体解包可以得到 Rc ,进而得到TreeNode 进行访问值和递归。

除了match,rust还提供了便捷的模式匹配语法。下面对145. 二叉树的后序遍历 进行说明

impl Solution {
    pub fn postorder_traversal(root: Option>>) -> Vec {
        fn postorder_dfs(node: &Option>>, ans: &mut Vec) {
            if let Some(x) = node {
                postorder_dfs(&x.borrow().left, ans);
                postorder_dfs(&x.borrow().right, ans);
                ans.push(x.borrow().val);
            }
        }

        let mut ans = vec![];
        postorder_dfs(&root, &mut ans);
        ans
    }
}

使用 if let 方式进行模式匹配,忽略掉匹配失败的分支,这里是忽略了 None 变体的 &Option结构。推荐使用最后一种方式,简单明了。

二叉树的 dfs 的递归算法写法很简单。并且由于函数传递的是 &Option, 在递归调用过程中,也不需要引用多个所有权。对于dfs的非递归写法,则需要借助栈。此时就涉及到二叉树节点的多所有权。

迭代写法

先序遍历

python3 的算法模板:

class Solution:
    def preorderTraversal(self, root: TreeNode) -> List[int]:
        return list(self.preorder(root))

    def preorder(self, node: TreeNode):
        if node is None:
            return
        # 节点进栈
        stack = [node]
        while len(stack) > 0:
            # 出栈
            node = stack.pop()
            # 访问节点
            yield node.val
            # 右子树进栈
            if node.right is not None:
                stack.append(node.right)
            # 左子树进栈
            if node.left is not None:
                stack.append(node.left)

用 Rust 直接翻译算法代码如下:

impl Solution {
    pub fn preorder_traversal(root: Option>>) -> Vec {
        Solution::preorder_dfs(root)
    }

    fn preorder_dfs(root: Option>>) -> Vec {
        let mut ans = vec![];

        if root.is_none() {
            return ans;
        }
        // 进栈
        let mut stack = vec![root];
        while !stack.is_empty() {
            // stack pop的值会自动包装 Option,需要调用 flatten 打平
            let node = stack.pop().flatten().unwrap();
            // 通过 Rc 的borrow 获取 Ref 节点
            let node = node.borrow();
            // 访问节点值
            ans.push(node.val);
            // 右子树进栈
            if let Some(ref right) = node.right {
                stack.push(Some(right.clone()));
            }
            // 左子树进栈
            if let Some(ref left) = node.left {
                stack.push(Some(left.clone()));
            }
        }
        ans
    }
}

由于 Rust 有Option 类型,实际上进栈出栈也可以直接操作,即使遇到子树不存在的 None,pop出来的时候也可以通过模式匹配处理。代码会比较紧凑。

impl Solution {
    pub fn preorder_traversal(root: Option>>) -> Vec {
        Solution::preorder_dfs(root)
    }

    fn preorder_dfs(root: Option>>) -> Vec {
        let mut ans = vec![];
        // 节点进栈,即使 root是None,也没关系。
        let mut stack = vec![root];
        while !stack.is_empty() {
            // 节点出栈则进行模式匹配,过滤不存在的节点 None
            if let Some(node) = stack.pop().flatten() {
                ans.push(node.borrow().val);
                // 右子树进栈,None也没关系
                stack.push(node.borrow().right.clone());
                // 左子树进栈,None也没关系
                stack.push(node.borrow().left.clone());
            }
        }
        ans
    }
}

向量 vec pop的结构被包装了 option。在进行先序遍历的时候,我们可以保证入栈出栈的节点都有值,就像第一种遍历的方式,那么可以直接入栈 Rc 结构,而不是 Option结构。

    fn preorder_dfs(root: Option>>) -> Vec {
        let mut ans = vec![];

        if root.is_none() {
            return ans;
        }
        // root 节点不是 None, unwrap 安全
        let mut stack = vec![root.unwrap()];
        while !stack.is_empty() {
            if let Some(node) = stack.pop() {
                ans.push(node.borrow().val);

                if let Some(ref right) = node.borrow().right {
                    // 右子树 Rc 结构进栈,不需要包装一层 option
                    stack.push(right.clone());
                }
                if let Some(ref left) = node.borrow().left {
                     // 左子树 Rc 结构进栈
                    stack.push(left.clone());
                }
            }
        }
        ans
    }

上面是rust 实现先序遍历的几种写法,实际问题可以具体考虑。还有一种先序遍历的算法模板,即左子树先进栈,进栈前访问节点。如果是出栈的时候访问节点,算法则变成中序遍历。并且后序遍历的算法结构也可以统一起来。

下面是 python的算法模板:

class Solution:
    def preorderTraversal(self, root: TreeNode) -> List[int]:
        return list(self.preorder(root))

    def preorder(self, node: TreeNode):
        stack = []
        while True:
            # 左子树进栈
            while node is not None:
                # 访问节点
                yield node.val
                # 节点进栈
                stack.append(node)
                # 依次遍历左子树
                node = node.left
            # 栈空,退出
            if len(stack) <= 0:
                break
            # 出栈
            # yield node.val 如果在此时访问 节点,则是中序遍历
            node = stack.pop()
            # 控制权转移到右子树
            node = node.right

上面的算法模板,和递归的逻辑很相似。也就是先处理左子树,然后处理右子树。下面是 rust 代码的实现:

impl Solution {
    pub fn preorder_traversal(root: Option>>) -> Vec {
        Solution::preorder(root)
    }

    fn preorder(root: Option>>) -> Vec {
        let mut ans = vec![];
        let mut stack = vec![];
        let mut node = root;

        loop {
            // 通过模式匹配活动节点的 Rc 引用
            while let Some(x) = node {
                ans.push(x.borrow().val);
                // 节点进栈
                stack.push(x.clone());
                // 转移左子树
                node = x.borrow().left.clone();
            }
            if stack.is_empty() {
                break;
            }
            node = stack.pop().unwrap().borrow().right.clone();
        }
        ans
    }
}

代码和之前的先序遍历不太一样,node 会重新赋值,因此它不能借用。需要把所有权转移给 x,x是 Rc 结构。可以clone增加引用计数。本质上是让 node重新绑定了新的所有权。

掌握了上面几种遍历方式,套用最后一种遍历方法。下面是中序和后序实现。

中序遍历

impl Solution {
    pub fn inorder_traversal(root: Option>>) -> Vec {
        Solution::inorder(root)
    }

    fn inorder(root: Option>>) -> Vec {
        let mut ans = vec![];
        let mut stack = vec![];

        let mut node = root;
        loop {
            while let Some(x) = node {
                stack.push(x.clone());
                node = x.borrow().left.clone();
            }
            if stack.is_empty() {
                break;
            }
            if let Some(x) = stack.pop() {
                ans.push(x.borrow().val);
                node = x.borrow().right.clone();
            }
        }
        ans
    }
}

后序遍历

impl Solution {
    pub fn postorder_traversal(root: Option>>) -> Vec {
        Solution::postorder(root)
    }

    fn postorder(root: Option>>) -> Vec {

        let mut ans = vec![];
        let mut stack = vec![];
        let mut node = root;
        let mut visited = None;

        loop {
            while let Some(x) = node {
                stack.push(x.clone());
                node = x.borrow().left.clone();
            }
            if stack.is_empty() {
                break;
            }
            if stack.last().unwrap().borrow().right != visited {
                node = stack.last().unwrap().borrow().right.clone();
                visited = None
            } else {
                if let Some(x) = stack.pop(){
                    ans.push(x.borrow().val);
                    visited = Some(x.clone());
                }
            }
        }
        ans
    }
}

至此,二叉树的DFS遍历方法介绍完毕,下面再看下广度搜索层序遍历

广度优先

层序遍历如其字面意思,先遍历二叉树的每一层所有节点,然后再向下推进一层。与先序遍历类似,层序遍历只需要借助一个 队列 queue。

102. 二叉树层序遍历

impl Solution {
    pub fn level_order(root: Option>>) -> Vec> {
        let mut ans = vec![];

        let mut queue = VecDeque::new();
        queue.push_back(root);

        while !queue.is_empty() {
            let size = queue.len();
            let mut level = vec![];
            for _ in 0..size {
                if let Some(x) = queue.pop_front().flatten() {
                    let node = x.borrow();
                    level.push(node.val);
                    queue.push_back(node.left.clone());
                    queue.push_back(node.right.clone());
                }
            }
            if !level.is_empty() {
                ans.push(level);
            }
        }
        ans
    }
}

二叉树反序列化

二叉树的层序遍历有很多用法。Leetcode 的test case 使用的是二叉树的层序遍历的数组,进行重建二叉树。下面就介绍这个方法。

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

impl TreeNode {
    #[inline]
    pub fn new(val: i32) -> Self {
        TreeNode {
            val,
            left: None,
            right: None,
        }
    }

    pub fn get_height(root: &Option>>) -> i32 {
        fn dfs(root: &Option>>) -> i32 {
            match root {
                None => 0,
                Some(node) => {
                    let node = node.borrow();
                    1 + max(dfs(&node.left), dfs(&node.right))
                }
            }
        }
        dfs(&root)
    }
    
    // 通过数组反序列化生成一棵树
    pub fn create(nums: Vec>) -> Option>> {
        if nums.is_empty() {
            return None;
        }
        let size = nums.len();
        let mut index = 0;
        let root = Some(Rc::new(RefCell::new(Self::new(nums[0].unwrap()))));
        let mut queue = VecDeque::new();
        queue.push_back(root.clone());
        while !queue.is_empty() {
            let q_size = queue.len();
            for _i in 0..q_size {
                if let Some(x) = queue.pop_front().flatten() {
                    let mut node = x.borrow_mut();
                    let lseq = 2 * index + 1;
                    let rseq = 2 * index + 2;
                    if lseq < size && nums[lseq].is_some() {
                        node.left = Some(Rc::new(RefCell::new(Self::new(nums[lseq].unwrap()))));
                        queue.push_back(node.left.clone());
                    }

                    if rseq < size && nums[rseq].is_some() {
                        node.right = Some(Rc::new(RefCell::new(Self::new(nums[rseq].unwrap()))));
                        queue.push_back(node.right.clone());
                    }
                }
                index += 1;
            }
        }
        root
    }
    
    // 将一棵树序列化成一个数组
    pub fn literal(root: Option>>) -> Vec> {
        if root.is_none() {
            return vec![];
        }
        let mut ans = vec![];
        let mut queue = VecDeque::new();
        queue.push_back(root);
        while !queue.is_empty() {
            let qsize = queue.len();
            for _ in 0..qsize {
                match queue.pop_front().flatten() {
                    Some(x) => {
                        ans.push(Some(x.borrow().val));
                        queue.push_back(x.borrow().left.clone());
                        queue.push_back(x.borrow().right.clone());
                    }
                    None => ans.push(None),
                }
            }
        }
        let size = ans.len();
        for i in (0..size).rev() {
            if ans[i].is_none() {
                ans.pop();
            } else {
                break;
            }
        }
        ans
    }
}

二叉树树形打印

有了二叉树的反序列化,可以方便的构造test case。通常为了直观验证一棵树,可以把树的拓扑树形进行打印。打印二叉树需要使用二叉树树高,因为满二叉树的宽度等于 width = (1<

leetcode 655. 输出二叉树

pub fn print_tree(root: Option>>) -> String {
    // 二叉树高度
    let height = TreeNode::get_height(&root);
    // 满二叉树的宽度
    let width = (1 << height) - 1;
    let mut ans = vec![vec![" ".to_string(); width as usize]; height as usize];

    // dfs 搜索
    fn dfs(ans: &mut Vec>, node: &Option>>, deep: usize, lo: usize, hi: usize) {
        if let Some(x) = node {
            let node = x.borrow();
            let mid = lo + (hi - lo) / 2;
            ans[deep][mid] = x.borrow().val.to_string();
            dfs(ans, &node.left, deep + 1, lo, mid);
            dfs(ans, &node.right, deep + 1, mid + 1, hi);
        }
    }

    dfs(&mut ans, &root, 0usize, 0usize, width as usize);
    // 将所有字符连起来
    ans.iter().map(|x| x.concat()).collect::>().join("\n")
}

上述打印的树形二叉树并没有 path 的边。需要变可以自行添加。

总结

通过rust构造二叉树,需要使用到 rust的特性,如所有权,借用检查,内部可变性等概念。另外掌握二叉树的基本性质和遍历方法,可以解决很多 leetcode上的题目。最后,通过二叉树的反序列化和树形打印,对刷题和学习验证相关算法都很有帮助。

更多 Rust 与Leetcode可以参考这个文档

你可能感兴趣的:(Rust 二叉树实现与遍历)