明日の授業でAtCoderのTypical Stairsという問題に取り組んでもらう(動的計画法の回)。
問題ページはこちら
ちなみにこの演習は以下の本をもとに進める。
先週が再帰を使った全探索だったのでその応用という流れ。1年前の記憶はもうないので予習していたら。。。詰まった。昨年はどうやって教えたのだろう。。。再帰ではどうもうまく書けない。なぜか解答がなかった。なんでだっけ?書籍にない例題にしたんだっけ?記憶がない。。。まぁいい。
対象は2年生なので再帰が自由自在に使えるという状況ではもちろんないのでサンプルを動かしてもらいながらヒントを小出しにして自分でなんとか完成する経験を積んでもらう感じの演習。
まず、似たような問題としてA - Frog 1という問題を例に再帰のコードを示す。たぶん、上記の教科書のコードだったと思う。解説等は上記の本にあるので知りたい方は買って読んでください。
# 十分に大きい値 INF = 2 ** 60 # 入力 N = int(input()) h = list(map(int, input().split())) # 再帰関数 def rec(i): # 終了条件 if i == 0: return 0 # 答え result = INF # 頂点i-1から来た場合 if i - 1 >= 0: result = min(result, rec(i-1) + abs(h[i] - h[i-1])) # 頂点i-2から来た場合 if i - 2 >= 0: result = min(result, rec(i-2) + abs(h[i] - h[i-2])) # 答えを返す return result # 再帰的に求める print(rec(N - 1))
同じ計算をしているので計算を省略するためリスト変数dpを使用したのが以下。
import sys sys.setrecursionlimit(10**6) # 十分に大きい値 INF = 2 ** 60 # 入力 N = int(input()) h = list(map(int, input().split())) # メモ化用の配列 dp = [-1] * N # 再帰関数 def rec(i): # 終了条件 if i == 0: return 0 # 計算済みであれば、その値を返す if dp[i] != -1: return dp[i] # 答え result = INF # 頂点i-1から来た場合 if i - 1 >= 0: result = min(result, rec(i-1) + abs(h[i] - h[i-1])) # 頂点i-2から来た場合 if i - 2 >= 0: result = min(result, rec(i-2) + abs(h[i] - h[i-2])) # 答えをメモ化しながら返す dp[i] = result return result # 再帰的に求める print(rec(N - 1))
とりあえずA - Frog 1はこれでACになる。
さて、本題。
Typical Stairsはたどり着くことができる経路数を答える必要があるので再帰関数の中身がちょっと変わる。まず、dp変数はおいておいて単純に考える。問題はAtCoderを見てね。
まず、何も考えないですべてのパスの組み合わせを再帰関数で作るコードを考えてみる。上のコードを使うことを前提としている。VSCodeのデバッガで動きをみたりしてだいたい行けそうだなと思って書いてみた。問題文と逆だが結果的には同じになるはずなので0段目からではなくN段目からスタートして0段目までいけるか確認する。行けたら1を返す。それだけ。rec(N-1)をrec(N)にする必要があるのがちょっとミソ。
# 入力 N = int(input()) # 再帰関数 def rec(i): # 終了条件 if i == 0: return 1 # 答え result = 0 # 頂点i-1から来た場合 if i - 1 >= 0: result += rec(i-1) # 頂点i-2から来た場合 if i - 2 >= 0: result += rec(i-2) # 答えを返す return result # 再帰的に求める print(rec(N))
そして、濡れている段はいけないことにする処理を書いてみたのが以下。これでAtCoderにある入力例の1と2は結果が出る。ただ、入力例3はTLEっぽい挙動をする(つまりいつまでも完了しない)。
# 入力 N, M = map(int, input().split()) a = [] for i in range(M): a.append(int(input())) # 再帰関数 def rec(i): # 終了条件 if i == 0: return 1 # 濡れた段? for ai in a: if i == ai: return 0 # 答え result = 0 # 頂点i-1から来た場合 if i - 1 >= 0: result += rec(i-1) # 頂点i-2から来た場合 if i - 2 >= 0: result += rec(i-2) # 答えを返す return result # 再帰的に求める print(rec(N))
そこでdp変数を導入する。これで入力例3も答えが返ってきた。これであとはAtCoderに投稿すれば終わりじゃんと思ったが甘かった。。。色々とネットで調べるけどなぜか再帰を使った例は見つからない。困ったなぁ。。。昨年の授業資料をみるとどうやら再帰を使うのを諦めたような痕跡があった。情けない。。。
import sys sys.setrecursionlimit(10**6) # 入力 N, M = map(int, input().split()) # メモ化用の配列 dp = [-1] * (N + 1) #← N を N+1 a = [] for i in range(M): a.append(int(input())) # 再帰関数 def rec(i): # 終了条件 if i == 0: return 1 # 計算済みであれば、その値を返す if dp[i] != -1: return dp[i] # 濡れた段? for ai in a: if i == ai: return 0 # 答え result = 0 # 頂点i-1から来た場合 if i - 1 >= 0: result += rec(i-1) # 頂点i-2から来た場合 if i - 2 >= 0: result += rec(i-2) # 答えをメモ化しながら返す dp[i] = result return result # 再帰的に求める print(rec(N) % 1000000007)
色々やってもTLEにしかならなくて諦めてfor文のやリかたを真似して書いていて気付いた。そうか。。。dp変数に濡れた段のdp変数を0にしておけばいちいちチェックしなくていいのか。。。
試しに以下のコードに変更してみた。
import sys sys.setrecursionlimit(10**6) # 入力 N, M = map(int, input().split()) # メモ化用の配列 dp = [-1] * (N + 1) #← N を N+1 for i in range(M): dp[int(input())] = 0 # 再帰関数 def rec(i): # 終了条件 if i == 0: return 1 # 計算済みであれば、その値を返す if dp[i] != -1: return dp[i] # 答え result = 0 # 頂点i-1から来た場合 if i - 1 >= 0: result += rec(i-1) # 頂点i-2から来た場合 if i - 2 >= 0: result += rec(i-2) # 答えをメモ化しながら返す dp[i] = result return result # 再帰的に求める print(rec(N) % 1000000007)
これでACになった。。。なるほど。濡れた段が多くなると組合せ数が多くなりTLEになるようになっていた問題だったということでした。ちなみにネットの解答はもっとシンプルです。競技プログラミングは奥が深いですね。。。
おしまい。
