C++版

はじめに

---

binary search tree の，先頭からk番目の数値を答える問題．

Rust を勉強中であり，その一環として LeetCode の問題を解いてみた．

コード全体

---

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

impl Solution {
  pub fn kth_smallest(root: Option<Rc<RefCell<TreeNode>>>, k: i32) -> i32 {
    fn helper(node: Option<Rc<RefCell<TreeNode>>>, k: i32) -> (Option<i32>, i32) {
      match node {
        Some(node) => {
          let n = node.borrow();
          match helper(n.left.clone(), k) {
            (None, count) => {
              match count {
                1 => (Some(n.val), 0),
                _ => helper(n.right.clone(), count - 1)
              }
            },
            (res, _) => (res, 0)
          }
        },
        None => (None, k)
      }
    }

    helper(root, k).0.unwrap()
  }
}

解説

---

事前準備

---

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

impl Solution {

関数定義

---

  pub fn kth_smallest(root: Option<Rc<RefCell<TreeNode>>>, k: i32) -> i32 {

ヘルパ関数の定義

---

k番目を探すという条件のためには，再帰された関数でいくつのノードを探索したのかを知る必要がある．そのため，全体のkth_smallest()とは別の型の戻り値をもつ関数をヘルパ関数として定義し，探索に用いる．

戻り値のタプルのうち，１つ目は結果，２つ目はk - それまでにカウントしたnode数を表すものとする．

    fn helper(node: Option<Rc<RefCell<TreeNode>>>, k: i32) -> (Option<i32>, i32) {

引数の有無により分岐

---

Option<T>は，空(None)が渡されることもあるので，有無によって処理を分岐する．

      match node {

(node が)ある場合

---

左側に再帰し，結果が存在すればそれをそのまま返す．

存在しない場合，自身がk番目であれば自身の値を返す．

そうでない場合は，カウントを引き継ぎながら右側に再帰する．

        Some(node) => {
          let n = node.borrow();
          match helper(n.left.clone(), k) {
            (None, count) => {
              match count {
                1 => (Some(n.val), 0),
                _ => helper(n.right.clone(), count - 1)
              }
            },
            (res, _) => (res, 0)
          }
        },

(node が)ない場合

---

これは，空のleftやrightでhelper()が呼ばれた時に発生する．

関数は呼ばれているが，ノードをカウントするわけではないため，結果なしNoneのカウントダウンなしkを返す．

        None => (None, k)
      }
    }

ヘルパ関数を呼び出す

---

    helper(root, k).0.unwrap()
  }
}

おわりに

---

C++でいう Lambda 式を，通常の関数と同等の書き方で定義することが出来るのは，非常に便利だと感じる．

また，match構文やOptionがかなり便利であり，書き方に柔軟性を持たせることができる．

ポインタの扱いが難しく，所有権に関してはやはり慣れない概念だなと思う．しかし，このおかげでメモリ安全性を担保してくれているというので，言語に従う事による安心感はやはりある．

参考

---