ABC275 D問題の考え方 (Python/Ruby)
ABC275 D - Yet Another Recursive Function の考え方を書きました。
問題の概要
# Python
def f(k):
if k==0:
return 1
else:
return f(k//2) + f(k//3)
# Ruby
def f(k)
if k==0
1
else
f(k/2) + f(k/3)
end
end
について、 f(N)
を求めます。
与えられるNの範囲は、 0以上 かつ 10^18以下 です。
愚直な実装をすると何が起きるか
上のコードにNの入力とf(N)の出力を追加しました。
# Python
N = int(input())
def f(k):
if k==0:
return 1
else:
# "k//N" は、Nで割って小数点以下切り捨て の意味
return f(k//2) + f(k//3)
print(f(N))
# Ruby
N = gets.to_i
def f(k)
if k==0
1
else
f(k/2) + f(k/3)
end
end
puts f(N)
これで、正しい答えを求められるようになりました。
しかし、このまま提出すると TLE(Time Limit Exceeded / 実行時間制限超過) になってしまいます。 この問題の実行時間制限 “2秒以内” で処理が終わらない場合がある、ということです。
入出力は整数たった1つですから、重いのは(実行時間のほとんどは) f()
の部分と考えてよいでしょう。
このようなとき、まず”ループの実行回数”を考えると、高速化の手がかりになります。
しかし、このコードにはfor文がありません。こまった。
無駄な処理はどこか
実動作から”無駄”を探してみることにします。
何を計算中か、都度出力するようにしてみました↓
# Python
def f(k):
if k==0:
return 1
else:
print(f'f({k})を計算中...')
return f(k//2) + f(k//3)
# Ruby
def f(k)
if k==0
1
else
puts("f(#{k})を計算中...")
f(k/2) + f(k/3)
end
end
さらに、見やすさのため、再帰の深さに応じてインデントを追加しました↓
# Python
def f(k, depth=0):
if k==0:
return 1
else:
print(' '*depth, f'f({k})を計算中...')
return f(k//2, depth+1) + f(k//3, depth+1)
# Ruby
def f(k, depth=0)
if k==0
1
else
puts(' '*depth + "f(#{k})を計算中...")
f(k/2, depth+1) + f(k/3, depth+1)
end
end
Pythonインタプリタで動かしてみると、以下の出力が得られました。
>>> f(36)
f(36)を計算中...
f(18)を計算中...
f(9)を計算中...
f(4)を計算中...
f(2)を計算中...
f(1)を計算中...
f(1)を計算中...
f(3)を計算中...
f(1)を計算中...
f(1)を計算中...
f(6)を計算中...
f(3)を計算中...
f(1)を計算中...
f(1)を計算中...
f(2)を計算中...
f(1)を計算中...
f(12)を計算中...
f(6)を計算中...
f(3)を計算中...
f(1)を計算中...
f(1)を計算中...
f(2)を計算中...
f(1)を計算中...
f(4)を計算中...
f(2)を計算中...
f(1)を計算中...
f(1)を計算中...
f(1)
が繰り返し計算されています。
その理由が、 f(2)
や f(3)
の計算において必要となるから ということもわかりました。
さらによく見ると、 f(6)
と f(4)
の計算がそれぞれ2回ずつ行われています。
もし、この計算を手作業でやった場合、「f(6)は先ほど計算したから…」と、計算結果を再利用することでしょう。
このコードは計算結果の記憶と再利用をしないため、全く同じ計算を繰り返してしまっています。
N=36
なら実行時間への影響は誤差程度ですが、特にN
が大きい(10^18
に近い)場合には重複した計算の実行回数が膨大になってしまいます。
結果を”メモ”しよう
重複した計算を行わないようにするため、計算結果をメモして再利用できるようにします。
# Python
memo = {}
def f(k):
if k==0:
return 1
elif k in memo: # まずはメモにないかチェック
return memo[k] # あるなら利用する
else:
print(f'f({k})を計算中...')
val = f(k//2) + f(k//3)
memo[k] = val # メモに書き込む
return val
# Ruby
@memo = {}
def f(k)
if k==0
1
elsif @memo.has_key?(k) # まずはメモにないかチェック
@memo[k] # あるなら利用する
else
puts("f(#{k})を計算中...")
val = f(k/2) + f(k/3)
@memo[k] = val # メモに書き込む
val
end
end
表示整形用のインデントは削除しました。
Pythonインタプリタで動かしてみると、以下の出力が得られました。
>>> f(36)
f(36)を計算中...
f(18)を計算中...
f(9)を計算中...
f(4)を計算中...
f(2)を計算中...
f(1)を計算中...
f(3)を計算中...
f(6)を計算中...
f(12)を計算中...
先程の結果と比べて、明らかに計算回数が減っています。
同じ入力に対しては常に同じ出力であり、かつ、同じ入出力が何度も必要になる という場面では、一度計算した”入力と出力のペア”をメモすることで無駄な処理を省くことができます。
このテクニックを再帰関数に適用する手法は「メモ化再帰」と呼ばれます。
ということで、メモ化再帰の典型問題でした。
計算量の考察については、公式解説をご確認ください。
実装例
- Python 3.8.2 (実行時間: 27 ms)
- 組み込みの @functools.lru_cache を使うと簡単にメモ化を実現できます。
- Ruby 2.7.1 (実行時間: 59 ms)
||=
演算子を使うと、memoの参照と代入をシンプルに書けます。