ABC006 B問題 トリボナッチ数列をメモ化再帰で解く

ABC006のB問題を解説します。

atcoder.jp

問題の要約

3つ前までの数字を足した数列（トリボナッチ数列）の項目を求める問題です。
数字が大きくなるのでで割ったあまりが求められています。
となります。

難しいポイント

が制約です。再帰関数を書いてしまうとタイムアウトしますし、スタックオーバーフローします。
あと単純に足すと数字がオーバーフローします。

解き方の方針

4項目からn項目まで順番に答えを配列に記録しておきます。
足すときに数字が大きく、求められていることがで割ったあまりなので、足した後にで割ったあまりを記録しておきます。
足し算の場合、途中で何回も割ったあまりを求めても答えは変わりません。

@drkenさんの記事を紹介しておきます。

「1000000007 で割ったあまり」の求め方を総特集！ 〜 逆元から離散対数まで 〜 - Qiita

実装

n = int(input())
MOD = 10 ** 4 + 7

# 最初の3項は答えが決まっているので、すぐに答えをprintして終了する
if n in [1, 2]:
    print(0)
    exit()
if n == 3:
    print(1)
    exit()

tribo_memo = [0 for _ in range(n + 1)]
tribo_memo[3] = 1
# 4項からn項まで順番に答えを求めて配列に記録する
for i in range(4, n + 1):
    tribo_memo[i] = (tribo_memo[i - 1] + tribo_memo[i - 2] + tribo_memo[i - 3]) % MOD

print(tribo_memo[n])