「問題解決力を鍛える!アルゴリズムとデータ構造」とLeetCodeで学ぶ再帰と動的計画法(メモ化)

概要

『問題解決力を鍛える! アルゴリズムとデータ構造』を読んだ内容を活かせる例を見つけたので、その解法を記事にしてみました。

今回の問題は、LeetCodeの1137. N-th Tribonacci Numberです。

まずは自分で挑戦したいという方は『問題解決力を鍛える! アルゴリズムとデータ構造』の再帰の章を読み、挑戦した後にこの記事の続きを読んでください。

詳細

内容

今回解く問題は以下のとおりです。

T0 = 0, T1 = 1, T2 = 1
n >= 0 のとき、Tn+3 = Tn + Tn+1 + Tn+2
与えられた n に対して、Tn の値を返します。

単純な再帰

まず、再帰を使った初期の解法は以下のようになります。

class Solution {
    func tribonacci(_ n: Int) -> Int {
        if n == 0 {
            return 0
        } else if n == 1 || n == 2 {
            return 1
        }
        return tribonacci(n - 1) + tribonacci(n - 2) + tribonacci(n - 3)
    }
}

しかし、これでは n=36の場合に時間制限に引っかかります。
これは、再帰呼び出しが重複しており、計算量が爆発的に増加するためです。

ここで登場するのがメモ化です。
メモ化は動的計画法の一種であり、再帰的な問題を解く際に特に有効です。再帰呼び出しによる処理の重複箇所を保存しておき、再利用することで、処理速度を向上させることができます。
今回はトリボナッチ数列を解くのに使用していますが、他にもフィボナッチ数列や、より複雑な最適化問題や探索問題など、さまざまな問題に適用可能です。

そして、再帰処理の実行結果をメモを使用して保存するように修正したコードが以下のようになります。

class Solution {
    func tribonacci(_ n: Int) -> Int {
        var memo: [Int: Int] = [:]
        memo[0] = 0
        memo[1] = 1
        memo[2] = 1
        if n < 3 {
            return memo[n]!
        }
        for i in 3...n {
            memo[i] = (memo[i - 1] ?? 0) + (memo[i - 2] ?? 0) + (memo[i - 3] ?? 0)
        }
        return memo[n]!
    }
}

この解法では、計算済みの結果を memo に保存し、計算済みであれば関数を再呼び出すのではなく値を直接返します。これにより、重複計算を避け、大幅な効率化が達成されます。

以上となります。
実際に本を読んでみたけどいまいち理解しきれない時には実際に解いてみるとより知識が定着しやすくなると思うので、そういった方の助けになれば幸いです。