Vecからバイナリツリーを生成する・バイナリツリーをVecに変換する (LeetCode)

前提

LeetCodeでは、以下のようなバイナリツリー（Binary 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,
        }
    }
}

Vecからバイナリツリーを生成する

テストデータとして、Vecからこのバイナリツリーを作成するには、以下のような関数を用意します。なお、ここでは-1をnullとして扱っています。

fn to_tree(vals: Vec<i32>) -> Option<Rc<RefCell<TreeNode>>> {
    fn build_tree(vals: &Vec<i32>, index: usize) -> Option<Rc<RefCell<TreeNode>>> {
        if vals.len() <= index || vals[index] == -1 {
            return None;
        }
        let node = Rc::new(RefCell::new(TreeNode {
            val: vals[index],
            left: build_tree(vals, 2 * index + 1),
            right: build_tree(vals, 2 * index + 2),
        }));
        Some(node)
    }
    build_tree(&vals, 0)
}

例えば、以下のようなバイナリツリーを生成したい場合、次のようにto_tree関数を使います。

let vals = vec![1, 2, 3, 4, 5, 6, -1, -1, -1, 7, 8];
let root = to_tree(vals);

バイナリツリーをVecに変換する

上記とは逆に、バイナリツリーをVecに変換するには、以下のような関数を用意します。ここでも-1をnullとして扱います。

fn to_vec(root: Option<Rc<RefCell<TreeNode>>>) -> Vec<i32> {
    let mut stack = vec![(0, root.unwrap().clone())];
    let mut vals = Vec::new();
    while let Some((index, parent)) = stack.pop() {
        while vals.len() <= index {
            vals.push(-1);
        }
        vals[index] = parent.borrow().val;

        if let Some(left) = &parent.borrow().left {
            stack.push((index * 2 + 1, left.clone()));
        }
        if let Some(right) = &parent.borrow().right {
            stack.push((index * 2 + 2, right.clone()));
        }
    }
    vals
}

使用例は以下の通りです。

let root = ...;
let vals = to_vec(root);
println!("{:?}", vals); //=> [1, 2, 3, 4, 5, 6, -1, -1, -1, 7, 8]

環境情報

• rustc 1.84.1 (e71f9a9a9 2025-01-27)