はじめに
ABC 251 E 問題 Takahashi and Animals を解くために考えたこと、ACできるPython3(PyPy3)コードを紹介します。
E.Takahashi and Animals
問題ページ
難易度 ： 水色 1227
考察
$Ai$ すべてについて 払う,払わないを考えられれば簡単ですが、計算量が $2^N$ となるので間に合いません。
ここで、餌やりがどういうルールであるか理解するために可視化してみます。
つまり $\ i ≦ k ≦ N\ $ 番目までの動物への餌やりが完了している状態 からは、$i-1$ 番目の動物に餌を与えるために $A_{i-1}$ 払って $\ i-1 ≦ k ≦ N\ $ 番目までの動物への餌やりが完了した状態 に遷移するか、もしくは $A_{i-2}$ 払って $\ i-2 ≦ k ≦ N\ $ 番目までの動物への餌やりが完了した状態 へ遷移します。
$f(x)$ : $\ x ≦ k ≦ N\ $ 番目までの動物への餌やりを行うための最小値 を定義することで、この遷移は以下のように表せます。
\begin{equation}
f(x)=
\begin{cases}
min\ (\ f(x+1)\ ,\ f(x+2)\ ) +A_x & \text{
$1≦x<N-1$}\\
A_{N-1} & x=N-1\\
A_N & x=N
\end{cases}
\end{equation}
つまり、$\ f(x)\ $ は $N$ から降順に決定していき、最終的に定まる $\ f(1)\ $ で 求める最小金額を得られそうです。また一度決定した $f(x)$ を記憶していおけば、$O(1)$ で $\ f(x-1)\ $ ,$\ f(x-2)\ $を求めることができるので、全体の計算量は $O(N)$ となり十分高速です。
さて、$f(x)$ を求める手段としては 動的計画法 (DP) も考えられますが、今回は メモ化再帰 でこの $f(x)$ を求め、記憶することにします。
注意点 f(1) が最適解とは限らない
$f(1)$ は $\ A_1\ $ 払う場合の最小金額ではありますが、そもそも 動物すべてに餌を与えるために必ずしも $\ A_1\ $ 払う必要はありません。
例えば $\ A_N +\ A_2 +\ ....\ $ と払うこともでき、こちらの方が最適な可能性も十分考えられます。
そこで、$f(x)$ を定義しなおします。
$f_{use}(x)$ : $A_1\ $ 払う 場合において$\ x ≦ k ≦ N\ $ 番目までの動物への餌やりを行うための最小値
\begin{equation}
f_{use}(x)=
\begin{cases}
min\ (\ f_{use}(x+1)\ ,\ f_{use}(x+2)\ ) +A_x & \text{
$1≦x<N-1$}\\
A_{N-1} & x=N-1\\
A_N & x=N
\end{cases}
\end{equation}
$f_{unuse}(x)$ : $\ A_1\ $ 払わない ($\ A_N\ $ 必ず払う) 場合において$\ x ≦ k ≦ N\ $ 番目までの動物への餌やりを行うための最小値
\begin{equation}
f_{unuse}(x)=
\begin{cases}
min\ (\ f_{unuse}(x+1)\ ,\ f_{unuse}(x+2)\ ) +A_x & \text{
$2≦x<N-1$}\\
A_{N-1}+A_N & x=N-1\\
A_N & x=N
\end{cases}
\end{equation}
これらから得られる $min(\ f_{use}(1)\ , \ f_{unuse}(2)\ )$ によって、最適解が求まります。
コード
python3
762 ms/ 2000 ms
from functools import lru_cache
# A0 払う場合に利用
@lru_cache(maxsize=None)
def f_use(i):
ans=0
if i==N-1:
return A[-1]
elif i==N-2:
return A[-2]
else:
use=f_use(i+1)+A[i]
unuse=f_use(i+2)+A[i]
ans+=min(use,unuse)
return ans
# A0 払わない場合に利用
@lru_cache(maxsize=None)
def f_ununse(i):
ans=0
if i==N-1:
return A[-1]
elif i==N-2:
return A[-1]+A[-2]
else:
use=f_ununse(i+1)+A[i]
unuse=f_ununse(i+2)+A[i]
ans+=min(use,unuse)
return ans
if __name__ == "__main__":
# 再帰上限設定
import sys
sys.setrecursionlimit(10**7)
N=int(input())
A=[int(a) for a in input().split()]
if N==2:
print(min(A[0],A[1]))
exit()
# A0 払う場合
use=f_use(0)
# A0 払わない場合
unuse=f_ununse(1)
print(min(use,unuse))
別解 動的計画法 (DP)
本方針と同じ考察で
➀ 前から何番目までの動物に餌を与えたか
➁ i 番目までの動物すべてに餌を与える場合の最小値
以上、2つの情報で状態を定義することで、遷移に忠実に状態を更新できるとわかります。(便宜上、遷移方向を本方針と逆にしたが本質は同じ)
1次元 $\ dp\ $ テーブルを作成し $dp[➀] = ➁$ とする場合、サイズは $N$ , 遷移数 2 で全体の計算量は $O(N)$ となり十分高速です。
こちらも本方針と同様、$A_1$ 払う場合 : $\ dp[0][i]\ $ と払わない場合 : $\ dp[1][i]\ $ に分けて考えます。
最終的には $A_1$ 払う場合、$\ dp[0][N-1]\ $ もしくは $\ dp[0][N]\ $ で最小金額が得られ、$A_1$ 払わない場合、$\ dp[1][N]\ $ で最小金額が得られます。
実装では余計な条件設定を省くために、テーブルサイズを拡大しています。
コード
pypy3
206 ms/ 2000 ms
N=int(input())
A=[int(a) for a in input().split()]+[10**20]*10
dp=[[10**20 for _ in range(N+10)] for _ in range(2)]
dp[0][0]=A[0]
dp[0][1]=A[0]
dp[1][0]=0
dp[1][1]=A[1]
for i in range(N):
dp[0][i+1]=min(dp[0][i+1],dp[0][i]+A[i])
dp[0][i+2]=min(dp[0][i+2],dp[0][i]+A[i+1])
dp[1][i+1]=min(dp[1][i+1],dp[1][i]+A[i])
dp[1][i+2]=min(dp[1][i+2],dp[1][i]+A[i+1])
ans=min(min(dp[0][N-1],dp[0][N]),dp[1][N])
print(ans)
補足
-
メモ化再帰
functools --- 高階関数と呼び出し可能オブジェクトの操作
デコレータ @lru_cache() を関数の上につけるだけで、その関数の返り値を記憶し、同じ引数で呼ばれた時に再計算を行うことなく帰り値が得られるようになります。
maxsize 引数で記憶する量の上限を設定することが可能です。デフォルトのままでは、128種以上記憶することしかできず、次々に記憶した値を忘れていくので再計算が発生する恐れがあります。そこで今回は None で上限を設定することなく全て記憶させるようにしました。 -
メモ化再帰類題
問題ページ
ネタバレが気に入らない方のために、abc 270 ～ 275 の C,D 問題 のどれかであることを記します。
こちらの問題についても記事をあげているので、ぜひご覧ください -
再帰上限設定
Pythonの再帰回数の上限を確認・変更（sys.setrecursionlimitなど
再帰回数には上限があり、それを超える場合には上限を変更する必要があります。
終わり
質問やご指摘はこちらまで
Twitter : Waaa1471