firi.のmemo

yukicoderの過去問 (yukicoder April Fool Contest 2019 G)

形式的冪級数解。

$O(K \log K)$ 解法(想定解は $O(K^2)$ を定数倍高速化しています)

まずは2次元DPを考える。 $dp[i][j]=i$ 回移動して $j$ の位置にいるときの組み合わせの数　とすると、答えは、 $\displaystyle \sum_{i=0}^{\infty}dp[i][k]$ 。各移動で最低1マスは進むため $i \le k$ だから、そこで計算を打ち切れば計算量は $O(K^2N)$ だが、当然TLEする。

配るDPを考えると、遷移は、 $dp[i][j]をdp[i+1][j+x_1],dp[i+1][j+x_2] \cdots, dp[i+1][j+x_N]$ に足すという操作。

これをFFTを使って高速化したい。そのため、DP配列を多項式と見ることにする。 $f(x)=x^{x_1}+x^{x_2}+\cdots x^{x_n}$ とすると、 $dp_{i+1}(x) = dp_i(x)*f(x)$ となる。さらに、 $dp_0(x)=1$ なので、 $dp_i(x) = {f(x)}^i$ がわかる。

このようにDPを多項式に置き換えていくと、解は、 $\displaystyle \sum_{i=0}^{\infty}f(x)^i$ の $x^k$ の係数である。

ここで、 $\displaystyle \sum_{i=0}^{\infty}f(x)^i = \frac{1}{1-f(x)}$ であるから $1-f(x)$ の逆数が求まれば良い。この解説のE問題のところにある方法を使えば、逆数は $O(K\log K)$ で求まるので、 $O(K \log K)$ でこの問題は解ける。問題はMODが $10^9+7$ であることだが、これはNTTを使ったり、bit分割するFFTをしたりすれば解決する。

コード

https://yukicoder.me/submissions/334144