「もっとRust脳を鍛える数学パズル」の試み。

「もっとプログラマ脳を鍛える数学パズル」をRustで書き直すのは、ボケ防止にちょうど良いかもしれない、と思った。

P.13 例題1 : メモ化と動的計画法

元のRubyコード(p.015)。

pre1_2.rb

M, N = 10, 100

@memo = {}
def check(remain, pre)
    return @memo[[remain, pre]] if @memo[[remain,pre]]

    return 0 if remain < 0
    return 1 if remain == 0

    cnt = 0
    pre.upto(M) do |i|
        cnt += check(remain - i, i)
    end
    @memo[[remain, pre]] = cnt
end

puts check(N, 2)

Rustベタ移植。

main.rs

use std::collections::HashMap;

fn main() {
    let mut memo: HashMap<(i64, i64), i64> = HashMap::new();
    println!("{}", check(&mut memo, 10, 100, 2));
}

fn check(
    memo: &mut HashMap<(i64, i64), i64>,
    max_seat_at_one_table: i64,
    remain: i64,
    pre: i64,
) -> i64 {
    match memo.get(&(remain, pre)) {
        Some(cnt) => return *cnt,
        _ => {
            if remain < 0 {
                return 0;
            } else if remain == 0 {
                return 1;
            } else {
                let mut count = 0;
                for i in pre..=max_seat_at_one_table {
                    count += check(memo, max_seat_at_one_table, remain - i, i);
                }
                memo.insert((remain, pre), count);
                return count;
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_check() {
        let max_seat_at_one_table = 10;
        let number_of_people = 100;
        let mut memo: HashMap<(i64, i64), i64> = HashMap::new();
        assert_eq!(
            check(&mut memo, max_seat_at_one_table, number_of_people, 2),
            437_420
        );
    }
}

Rubyのuptoと同様の動きをRustのforにやらせるには、+1しないといけない。
と思ってたら @scivola さんからご指摘。ありがとうございます!

Rubyのショートコードでは何とも思わないグローバル変数だが、さすがにRustでやると気持ち悪い、というかやり方が分からないので、持ち回っている。Amazonの書評にもあったけど、変数を数学っぽく省略すると訳が分からなくなるので、ひどいところだけ、何となく長い名前にした。

改めて、条件分岐の処理にRubyの気楽さ(=プログラマへの信頼)を感じる。Rustはコンパクトながら抜けを許さないところが、これはこれで良い感じ。

2020/09/29追記

グローバル変数の扱いの指摘を受けて再度実装。グローバル変数をstructとして実現する案。これはもうデザインパターンに片足突っ込んでますね。

main.rs

use std::collections::HashMap;

struct Checker {
    memo: HashMap<(i64, i64), i64>,
    max_seat_at_one_table: i64,
}

impl Checker {
    pub fn check(&mut self, remain: i64, pre: i64) -> i64 {
        match &self.memo.get(&(remain, pre)) {
            Some(cnt) => return **cnt,
            _ => {
                if remain < 0 {
                    return 0;
                } else if remain == 0 {
                    return 1;
                } else {
                    let mut count = 0;
                    for i in pre..=self.max_seat_at_one_table {
                        count += self.check(remain - i, i);
                    }
                    &self.memo.insert((remain, pre), count);
                    return count;
                }
            }
        }
    }
}

fn main() {
    let mut chk = Checker {
        memo: HashMap::new(),
        max_seat_at_one_table: 10,
    };
    println!("{}", chk.check(100, 2));
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_check() {
        let number_of_people = 100;
        let mut chk = Checker {
            memo: HashMap::new(),
            max_seat_at_one_table: 10,
        };
        assert_eq!(chk.check(number_of_people, 2), 437_420);
    }
}