階段ピョンピョン1・2・3！【解説】

• 本稿はCodeIQで2014年2月3日～2月24日・2015年1月19日～3月2日に出題された 「コード銀行：階段ピョンピョン1・2・3！」 という問題の解答例と簡単な解説です。

• 本稿におけるコードは単なる解答の1例に過ぎません。「もっと良い解法があるよ！」とか、「こう書いた方がわかりやすいよ！」とか、「他の言語で書いてみたよ！」という場合は、自由にコメントに書いていただいても、個人のブログ等で公開していただいても結構です。

問題本文はこちら

階段ピョンピョン1・2・3！【解説】

◆基本的な考え方

n段目(n>3)に到達するとき、

[1] (n-1)段目まで上ったあと1段上がる
[2] (n-2)段目まで上ったあと2段上がる
[3] (n-3)段目まで上ったあと3段上がる

の3通りあるので、これらを再帰的に足し合わせるコードを書きます。

def count(n)
  return count(n-1) + count(n-2) + count(n-3)
end

という感じですね。もちろんnが負になってしまってはいけませんので、

if n < 0
  return 0

みたいな処理は必要です。nがちょうど0の場合は少し考える必要があります。

◆n = 0の場合の定義

たとえばnが2の場合を考えると、

[1] (2-1=1)段目まで上ったあと1段上がる
[2] (2-2=0)段目まで上ったあと2段上がる
[3] (2-3=-1)段目まで上ったあと3段上がる

[1]は明らかに1通りです。[3]は段数が負になるので0通りとしましょう。2段の上がり方は、実際には

・1段 + 1段
・一気に2段

の2通りです。 したがって、[2]は1通りと数えた方が都合が良くなりますので、count(0)は1と定義しておきましょう。

◆メモ化

単純な再帰の場合は計算量がO(n³)でかなり遅いので、メモ化を行います。

def count(n)
  if memo[n]が存在
    return memo[n]
  endif
  
  memo[n] = count(n-1) + count(n-2) + count(n-3)
  return memo[n]
end

のような感じですね。memo[n]が存在するかどうかの判定には、memo[n]を0で初期化しておいて、0より大きいときは存在するということにしておきましょう。つまり、

memo = [1, 0, 0, 0, ... , 0 ]

のような初期値を持ったリストを用意するということです。先程count(0)は1と定義しましたので、memo[0]の値のみ1で初期化しておきます。

◆コード例

以上を踏まえて、Python3でコードの例を書いておきました。

import sys
from operator import add
from functools import reduce

n = int(sys.argv[1]) # 段数はARGVから取る
memo = [1] + [0] * n # memo[0]は1で初期化、memo[1..n]は0で初期化
step = [1, 2, 3] # ステップは3種類

def count(n):
    if memo[n] > 0: return memo[n]
    memo[n] = reduce(add, [count(n-k) for k in step if n-k >= 0])
    return memo[n]

print(count(n))

n=30で実行すると、

53798080

という出力が得られます。