概要
あまり見直していないので細かい数式の誤植が心配ですが方針は理解できると思います.
解法
方針
目的を達成したあとも永遠にサイコロを振り続けることを考えて,次のように形式的べき級数 $P(x),F_i(x)$ を定めます.
- $P(x) = \sum p_nx^n$, $p_n$ は $n$ 回目に初めて $X=0$ を達成する確率.
- $F_i(x) = \sum b_nx^n$, $b_n$ は $X=i$ からスタートしたときに $n$ 回目に $X=0$ を達成する確率.
$F_R(x) = P(x)F_0(x)$ が成り立ちます.求めるものは $\sum_i ip_i = P'(1)$ です.
$G_i(x) = (1-x)F_i(x)$ とします.$P(x)=\frac{G_R(x)}{G_0(x)}$ より答は
$P'(1) = \dfrac{G_R'(1)G_0(1)-G_R(1)G_0′(1)}{G_0(1)^2}$
です.なお,$1-x$ 倍により $x=1$ が代入できる有理式になっていることなどが後述の $G_i$ の計算の結果として分かります.この辺はこの問題の私の解説でも同じ手法を使っているため深くは触れません.
$F_i(x)$, $G_i(x)$ の計算
まずは $F_i(x)$ を理解することにします.$f(x) = \frac{1}{K}(x+x^2+\cdots+x^K)$ とします.
$X=i$ からはじめてちょうど $n$ 回で目的を達成する確率は,$x^{-i}f(x)^n$ において指数が $M$ の倍数であるような部分の係数和に等しいです.
そこで,$\mu_M\subset \mathbb{C}$ を $1$ の $M$ 乗根全体とします.すると,$X=i$ からはじめてちょうど $n$ 回で目的を達成する確率は
$$\frac{1}{M}\sum_{\zeta\in \mu_M}\zeta^{-i}f(\zeta)^n$$
となります.ここから
$$F_i(x)=\frac{1}{M}\sum_{\zeta\in \mu_M}\frac{\zeta^{-i}}{1-f(\zeta)x}$$
が分かります.$f(1)=1$ より
$$M\cdot G_i(x)=1+\sum_{\zeta\in \mu_M\setminus\{1\}}\frac{\zeta^{-i}(1-x)}{1-f(\zeta)x}$$
となります.
このことを用いて答の計算式を考えます.まず $G_i(1)=1$ です.よって答は単に $G’_R(1)-G’_0 (1)$ となり,$G’_i(1)$ の計算に帰着されます.
$\frac{1-x}{(1-ax)}$ を微分して $x=1$ を代入すると $\frac{1}{a-1}$ になることから,
$$G_i'(1)=\sum_{\zeta\in \mu_M\setminus\{1\}}\frac{\zeta^{-i}}{f(\zeta)-1}$$
となります.よって,「$\zeta$ の有理式の $\zeta\in \mu_M$ に関する総和を求めよ」という種類の計算ができればよいことが分かります.
分母の有理化
$\zeta\in \mu_M$ は $1+x+\cdots+x^{M-1}=0$ の解です.
$a(x)=1+x+\cdots+x^{M-1}$, $b(x)=f(x)-1$ とすると,$a,b$ は互いに素な多項式です.したがって $a(x)p(x)+b(x)q(x)=1$ を満たす多項式 $p,q$ が存在します($q$ は $M-1$ 次未満にとれます).これらを用いると
$$\frac{\zeta^{-i}}{f(\zeta)-1}=\frac{\zeta^{-i}}{b(\zeta)}=\frac{\zeta^{-i}q(\zeta)}{b(\zeta)q(\zeta)}=\frac{\zeta^{-i}q(\zeta)}{1-a(\zeta)p(\zeta)}=\zeta^{-i}q(\zeta)$$
より,有理式を多項式に変換できます.
$$G_i'(1)=\sum_{\zeta\in \mu_M\setminus\{1\}}\zeta^{-i}q(\zeta)=-q(1)+\sum_{\zeta\in \mu_M}\zeta^{-i}q(\zeta)$$
ですが,$\sum_{\zeta\in \mu_M}\zeta^k=\begin{cases}M&(M\mid k)\\0& (M\nmid k)\end{cases}$ となるため,
$$G_i'(1)=\sum_{\zeta\in \mu_M\setminus\{1\}}\zeta^{-i}p(\zeta)=-q(1)+M[x^i]q(x)$$
となることが分かりました.ここから計算すべきは $[x^i]q(x)-[x^0]q(x)$ となり,結局 $q$ の係数の計算に帰着できました.
なお $q$ は $M-1$ 次未満にとれましたが,$i=M-1$ に対する答が $\mod 998244353$ で定義可能という制約から $[x^0]q(x)$ が $\mod 998244353$ で定義可能と分かり,さらにすべての $i$ に対する制約からすべての係数が $\mod 998244353$ で定義可能と分かります.よって以下の計算をすべて $\mod 998244353$ で行えることが分かります.
$p$ の係数の計算
この辺から大体公式解説と合流していると思います.大体次のような問題を $\mathbb{F}_{998244353}$ で解くことに帰着されています.
- $a(x)=1+x+\cdots+x^{M-1}$ とする.
- $b(x)=(x+x^2+\cdots+x^K)/K-1$ とする.
- $a(x)p(x)+b(x)q(x)=1$ となる最小次数の多項式 $q$ に対して $[x^i]q$ を計算せよ.
単純に多項式 GCD をするには $a(x)$ の次数が大きすぎるのでこの点に対処します.
$a$ を $b$ で割ったときの商を $A(x)$,余りを $r(x)$ として $a(x)=A(x)b(x)+r(x)$ と書きます.すると $r(x)p(x)\equiv 1\pmod{b(x)}$ となるので,まず $p(x)$ が $O(K\log^2K)$ 時間で計算できます(https://judge.yosupo.jp/problem/inv_of_polynomials).
すると求めるべきは $q(x)=\dfrac{1-a(x)p(x)}{b(x)}$ ですが,$(1-x)a(x)$ が疎であることに注目してこれを
$$q(x)=\frac{(1-x)-\bigl((1-x)a(x)\bigr)p(x)}{(1-x)b(x)}$$
などと書けば,疎な多項式と低次の多項式しか登場しない形になり,有理式で表される形式的べき級数の係数の計算に帰着できて答を $O(K\log K\log M)$ 時間で求めることができます.
全体の計算量は $O(K\log^2K+K\log K\log M)$ です.