LeetCodeに出てくるTreeNodeを簡単に作る in Rust

2022.08.08

TECH

LeetCode に出てくる TreeNode を簡単に作る in Rust

LeetCode の Tree 系の問題に以下のような木構造の構造体が出てきます。

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

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

これを用いて、二分木などの問題を解くように求められる訳ですが、
ローカルで問題を解こうと思った場合に、問題の通りの木構造を Rust で作るのが面倒でしたので、
これを文字列から作成できるようにしました。

拙いですが、以下がそのコードです。

use core::fmt;
use std::cell::RefCell;
use std::collections::VecDeque;
use std::error::Error;
use std::rc::Rc;

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

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

#[derive(Debug, Clone)]
pub struct TreeNodeError;
impl TreeNodeError {
    pub fn new() -> Self {
        Self
    }
}

impl fmt::Display for TreeNodeError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(&self.description())
    }
}

impl Error for TreeNodeError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        None
    }

    fn description(&self) -> &str {
        "Could not construct a TreeNode"
    }

    fn cause(&self) -> Option<&dyn Error> {
        self.source()
    }
}

impl TryFrom<&str> for TreeNode {
    type Error = Box<dyn Error>;

    fn try_from(value: &str) -> Result<Self, Self::Error> {
        let value = value.replace('[', "").replace(']', "");
        let value = value.split(',').map(|x| x.trim()).collect::<Vec<&str>>();
        let mut index = 0;
        let mut temp = VecDeque::new();
        let mut root: Option<Rc<RefCell<TreeNode>>> = None;

        if value.is_empty() {
            return Err(Box::new(TreeNodeError::new()));
        }

        loop {
            if root.is_none() {
                let val = value[index].parse()?;
                index += 1;
                let val = Rc::new(RefCell::new(Self::new(val)));
                root = Some(Rc::clone(&val));
                temp.push_front(Some(val));
                continue;
            }

            if index < value.len() && temp.is_empty() {
                return Err(Box::new(TreeNodeError::new()));
            }

            if index + 1 >= value.len() {
                break;
            }

            if let Some(node) = temp.pop_back() {
                let target = node.ok_or_else(TreeNodeError::new)?;
                let mut target = target.borrow_mut();

                if let Ok(left) = value[index].parse() {
                    let left = Rc::new(RefCell::new(Self::new(left)));
                    target.left = Some(Rc::clone(&left));
                    temp.push_front(Some(left));
                }
                index += 1;

                if let Ok(right) = value[index].parse() {
                    let right = Rc::new(RefCell::new(Self::new(right)));
                    target.right = Some(Rc::clone(&right));
                    temp.push_front(Some(right));
                }
                index += 1;
            }
        }

        let root = root.ok_or_else(TreeNodeError::new)?;
        let root = root.borrow();
        let root = root.clone();
        Ok(root)
    }
}

こちらの Playground で試せます。

エラーハンドリングはいい加減です。正しい構造の配列は少なくともきちんと処理できるかと思います。

これを用いると、

このような感じで作成していた TreeNode を

let root = TreeNode::new(5);
root.left = Some(Rc::new(RefCell::new(TreeNode::new(4))));
root.right = Some(Rc::new(RefCell::new(TreeNode::new(8))));
if let Some(node) = root.left.as_ref() {
    node.borrow_mut().left = Some(Rc::new(RefCell::new(TreeNode::new(11))));
    if let Some(node) = node.borrow_mut().left.as_ref() {
        node.borrow_mut().left = Some(Rc::new(RefCell::new(TreeNode::new(7))));
        node.borrow_mut().right = Some(Rc::new(RefCell::new(TreeNode::new(2))));
    }
}
if let Some(node) = test.right.as_ref() {
    node.borrow_mut().left = Some(Rc::new(RefCell::new(TreeNode::new(13))));
    node.borrow_mut().right = Some(Rc::new(RefCell::new(TreeNode::new(4))));
    if let Some(node) = node.borrow_mut().right.as_ref() {
        node.borrow_mut().right = Some(Rc::new(RefCell::new(TreeNode::new(1))));
    }
}

こう書けます。

let root = TreeNode::try_from("[5,4,8,11,null,13,4,7,2,null,null,null,1]").unwrap();

これで、LeetCode の Tree の問題をローカルで簡単に解けますね！

Tree を使う問題

以下のリンクに今回のコードを活かせる

https://leetcode.com/explore/learn/card/data-structure-tree/

ちなみに後で気がついたのですが、上記の問題集の最後に、このようなデシリアライズ処理自体を書かせる問題もありました。

まとめ

Rust 勉強中のため色々アルゴリズムの問題を解いてみているのですが、

今回はその中でツール的に作成した、Tree 構造を文字列から作成するコードを紹介しました。

小ネタが増えていますが、細々とでも更新継続していきます。

Homeへ戻る