AGC044 の A 問題「Pay to Win」を Python で解いてみる

再帰を書くのが苦手です.
問題はこちら：
atcoder.jp

問題概要

次の選択肢を繰り返すことで初期値 $0$ を $N$ にするための最小コストを求めよ.

コスト $A$ を払うことで, 数値を $2$ 倍する
コスト $B$ を払うことで, 数値を $3$ 倍する
コスト $C$ を払うことで, 数値を $5$ 倍する
コスト $D$ を払うことで, 数値を $1$ 増やす
コスト $D$ を払うことで, 数値を $1$ 減らす

ただし, テストケースは $T$ 個ある

制約

$\begin{align*} &1\le T \le 10 \\ &1\le N\le 10^{18}\\ &1\le A,B,C,D \le 10^9 \end{align*}$

$T$ は最大で $10$ なので, 気にしないことにします.

時間を逆転させて, 問題を言い換える

$N$ は非常に大きいので, すべての $n=1,2,\cdots, N$ に対して調べることはできません.
こういうときには逆向きに考えて得ができないかを考えるという定石があります.
今回の問題では $N$ から出発して $0$ になるための最小コストはいくらかを考えてみましょう.

現時点で枚数が $n$ のとき, 選択肢は

コスト $D$ で $n$ を1増やす
コスト $D$ で $n$ を1減らす
（ $2$ の倍数のときは）コスト $A$ で $n$ を $2$ で割る
（ $3$ の倍数のときは）コスト $B$ で $n$ を $3$ で割る
（ $5$ の倍数のときは）コスト $C$ で $n$ を $5$ で割る

の $2$ 個～ $5$ 個の選択肢を繰り返すことで, 初期状態 $n = N$ から終端状態 $n=0$ までにかかる最小コストを計算するということが考えられます.

しかし, $A, B, C$ が大きい場合には, $1$ 減らすを $O(N)$ やらないと最善がわからないため, 「コスト $D$ で $N$ を1増やす／減らす」をシミュレートすると間に合いません.

さらに言い換える

実は選択肢は

$n$ を $1$ 減らし続ける
$n$ を $2$ の倍数に増やしてから $2$ で割る
$n$ を $2$ の倍数に減らしてから $2$ で割る
$n$ を $3$ の倍数に増やしてから $3$ で割る
$n$ を $3$ の倍数に減らしてから $3$ で割る
$n$ を $5$ の倍数に増やしてから $5$ で割る
$n$ を $5$ の倍数に減らしてから $5$ で割る

の $7$ 通りの選択肢を繰り返すことで得られます. なぜならば, $q = 2,3,5$ で割って, 枚数 $n_1$ から枚数 $n_2$ への遷移を考える場合, なるべく早くに割ってから $\pm 1$ を行う方がコスト $D$ を払う回数が少なくて済むからです.

では, 遷移を数式に直すために, まずは「 $5$ の倍数に減らしてから $5$ で割る」をするのに掛かるコストを考えてみましょう.
これは
\begin{equation*}
D \times (n -n\text{ に近い } 5\text{ の倍数 }) + C
\end{equation*}
と表せます.

DPの漸化式をつくる

以上のことをほかの選択肢についてもまとめましょう.
$\text{dist}(n)$ を持っている数を $n$ にするための最小のコインの枚数とし, $\text{dist}(k)$ , $k=0,1,\cdots, n-1$ はすでに定まっているものとしましょう.
すると, $n \ge 2$ のとき, $\left\lceil \frac{n}{2}\right\rceil < n$ となることにより,
\begin{align*}
\textrm{dist}(n)
= \min\{
&D\times n,\\
& D\times \left(n - \left\lfloor \frac{n}{5}\right\rfloor *5\right) + C + \textrm{dist}\left( \left\lfloor \frac{n}{5}\right\rfloor\right),\\
& D\times \left( \left\lceil \frac{n}{5}\right\rceil *5 - n \right) + C + \textrm{dist}\left( \left\lceil \frac{n}{5}\right\rceil\right),\\
& D\times \left(n - \left\lfloor \frac{n}{3}\right\rfloor *3\right) + B + \textrm{dist}\left( \left\lfloor \frac{n}{3}\right\rfloor\right),\\
& D\times \left( \left\lceil \frac{n}{3}\right\rceil *3 - n \right) + B + \textrm{dist}\left( \left\lceil \frac{n}{3}\right\rceil\right),\\
& D\times \left(n - \left\lfloor \frac{n}{2}\right\rfloor *2\right) + A + \textrm{dist}\left( \left\lfloor \frac{n}{2}\right\rfloor\right),\\
& D\times \left( \left\lceil \frac{n}{2}\right\rceil *2 - n \right) + A + \textrm{dist}\left( \left\lceil \frac{n}{2}\right\rceil\right)\\
\}
\end{align*}
が成り立つことがわかります.
これに, 初期条件 $\text{dist}(0)=0$ , $\text{dist}(1)=D$ を加えれば, 動的計画法により問題を解くことができます.

def dist(n):
    if n == 0:
        return 0
    if n == 1:
        return D
    
    ret = min(
        D * n,
        D * abs(n - n//5*5) + C + dist(n//5),
        D * abs(n - (n+4)//5*5) + C + dist((n+4)//5),
        D * abs(n - n//3*3) + B + dist(n//3),
        D * abs(n - (n+2)//3*3) + B + dist((n+2)//3),
        D * abs(n - n//2*2) + A + dist(n//2),
        D * abs(n - (n+1)//2*2) + A + dist((n+1)//2)
    )

    return ret

また, 一度計算した結果は使いまわせるように, memo化しておきましょう.
これでACできます.
Submission #13550585 - AtCoder Grand Contest 044