概要
形式的べき級数による解の表示
とりあえず、次の言い換えができます。
- gold ならば $N$ 点得る。
- silver ならば $1$ 点得る。
- rainbow ならば $KN$ 点得る。
- $MN$ 点以上を得た瞬間にゴール
- ゴールまでにかかる回数の期待値を求めよ。
遷移の確率に合わせて、多項式 $f$ を $f = (1-p-q-r) + qx + px^N + rx^{NK}$ により定めます。
「はじめて $MN$ 点以上になるまでの回数」を言い換えます。$k$ 回目の出発時点でちょうど $n$ 点である確率を $p_{k,n}$ とします。すると、答は $\sum_{k=0}^{\infty}\sum_{n=0}^{NM-1} p_{k,n}$ と書き換えられます。
$p_{k,n} = [x^n] f^k$ より $\sum_{k=0}^{\infty} p_{k,n} = [x^n] \frac{1}{1-f}$ です。したがって、答は $\sum_{n=0}^{NM-1} [x^n] \frac{1}{1-f}$ に等しいことが分かります。
これは、$[x^{NM-1}]\frac{1}{(1-f)(1-x)}$ と等しいです。よって、問題は次のように書き替えられます。
$f = (p+q+r) – qx – px^N – rx^{NK}$ とするとき、$[x^{NM-1}]\frac{1}{f(1-x)}$ を答えよ。
計算誤差との戦い
これを FFT に基づく $O(NK\log NK \log NM)$ 時間アルゴリズムで計算すると、私の実装では計算誤差が出まくって、WA だらけになってしまいました。FFT をよりナイーブな多項式乗算に書き換えても同じことでした。線形漸化式の第 $N$ 項を計算するアルゴリズムの実装方法はいくつかあるので、実装によってはこの方法でも上手くいくのかもしれません。
（既存 AC コード、1 件しかないので比べようがないですね。出題者解説では FFT で計算量が落ちるようなことが書いてあるが、大丈夫なのだろうか？64bit より高い精度の小数を使うことは、私は試していません。）
とりあえず、主要部を分けて正確に計算することにして、計算誤差を減らします。$f(x)$ は $x = 1$ に $1$ 位の零点を持ちます。$2$ 位の零でないことは、$f'(1) \neq 0$ から分かります。
よって部分分数分解より、定数 $c$ を用いて
$$\frac{1}{f(x)(1-x)} = \frac{c}{(1-x)^2} + \frac{g(x)}{f(x)}$$
と表すことができます。$c$ は $\lim_{x\to 1}\frac{1-x}{f(x)} = \frac{1}{f'(1)}$ により計算できて、$g(x)$ は $g(x) = (1-x)^{-1}(1 – c\frac{f(x)}{1-x})$ により計算できます（$1-x$ で割ることは係数の累積和と対応することを利用して計算します）。
$[x^D] \frac{1}{f(x)(1-x)} = c(D+1) + [x^D]\frac{g(x)}{f(x)}$ となり、あとは分母が $f$ の有理式の係数を評価すればよくなりました。
この書き変えは地味ですが、$f(x) = (p+q+r) \biggl(1 – \frac{q}{p+q+r}x – \frac{p}{p+q+r}x^N – \frac{r}{p+q+r}x^{KN}\biggr)$ は、定数項以外の係数が非負です。このことは、$\text{rev}(f)$ による除算が加法の繰り返しのみで行えることを意味しています。
したがって、$x^N$ を $\text{rev}(f)$ で割った余りを求める計算を、非負実数の加法のみで実現することができ、桁落ちによる計算誤差が生じにくくなります。この方法で AC がとれました。
多項式乗算に FFT を利用するとやはり計算誤差で WA になってしまいましたが、上手くいくのでしょうか？提出が少ないのでよくわかりませんでした。いかがでしたか？
