等比数列に単項式が掛かっているものの和を計算するyosupo judgeの問題

スナネコです。yosupo judgeにある $\sum_{i=0}^{n-1}r^i i^d$ を解きました。
https://judge.yosupo.jp/problem/sum_of_exponential_times_polynomial

yosupo judge の github を見ると参考文献が載ってます。
https://github.com/yosupo06/library-checker-problems/issues/98

そういうわけで、この記事は Min_25 さんによる以下の記事を私の理解で書き直したものです。
ひとによってはオリジナルの方がわかりやすいかもしれません。
https://web.archive.org/web/20170622095007/http://min-25.hatenablog.com/entry/2015/04/24/031413

話の流れはこうです。
・数列 $a_n=\sum_{i=0}^{n-1}r^i i^d$ を漸化式で表す
・係数行列の固有多項式を求める
・固有多項式で割ったあまりを explicit な式で表す

ステップ1. 漸化式を出す

本題に入る前に、単項式 $n^d$ の値を $n$ に関する漸化式で求めるにはどうすればよいか考えてみましょう。
全ての $i\leq d$ についての $n^i$ を持てばできますね。具体的には、 $(n+1)^d$ が $n^i$ と二項係数を使って書けることから言えます。例えば
$\pmatrix{ 1 & 0 & 0 & 0\\ 1 & 1 & 0 & 0\\ 1 & 2 & 1 & 0\\ 1 & 3 & 3 & 1\\ } \pmatrix{ 1\\ n\\ n^2\\ n^3\\ }= \pmatrix{ 1\\ (n+1)\\ (n+1)^2\\ (n+1)^3\\ }$
となります。二項係数を並べたこの下三角行列のことを下三角パスカル行列というそうです。 $n\times n$ の下三角パスカル行列をwikipediaでは $L_n$ と表しているので、ここでもそう書くことにします。

本題に戻ります。ここまでわかれば元の問題を漸化式で書くのは簡単です。
$\pmatrix{ 1r & 0r & 0r & 0r & 0\\ 1r & 1r & 0r & 0r & 0\\ 1r & 2r & 1r & 0r & 0\\ 1r & 3r & 3r & 1r & 0\\ 0 & 0 & 0 & 1 & 1 } \pmatrix{ r^n\\ r^n n\\ r^n n^2\\ r^n n^3\\ \sum_{i=0}^{n-1} r^i i^3 }= \pmatrix{ r^{n+1}\\ r^{n+1} (n+1)\\ r^{n+1} (n+1)^2\\ r^{n+1} (n+1)^3\\ \sum_{i=0}^{(n+1)-1} r^i i^3 }$
こうですね。左上の $(d+1)\times(d+1)$ の小行列を $L_{d+1}$ とまとめて書くことにして
$M:=\pmatrix{ & 0\\ rL_{d+1} & \vdots\\ & 0\\ 0\ 0\ \ldots \ 0\ 1 & 1 }$
$A_n:=\pmatrix{ r^n\\ r^n n\\ \vdots\\ r^n n^d\\ \sum_{i=0}^{n-1} r^i i^d }$
とします。求めたいものは $A_n$ の最後の成分であり、 $A_n=M^nA_0$ 、 $A_0=(1,0,0,\ldots,0)$ です。
ここまでくれば、線形漸化式を持つ数列の第 $n$ 項を求めるのと同じようにできそうな気がしてきましたね。

ステップ2. 係数行列の固有多項式を求める

なぜ固有多項式を求めたいのか一応説明しておきましょうか。一般に、行列 $A$ の固有多項式 $p_A$ は $p_A(A)=O$ を満たします。このことから、 $f(x)=(x^n \bmod p(x))$ とすると、 $M^n=f(M)$ となります。*1
そうすると、小さい次数の項しか残らないので、いろいろ話が簡単になるわけです。

それでは本題に戻ります。さっき求めた係数行列 $M$ の固有多項式 $p$ を求めましょう。
$M$ は下三角行列で、対角成分は $r$ が $d+1$ 個、 $1$ が $1$ 個なので、 $p(x)=(x-r)^{d+1}(x-1)$ です。……はい、これだけです。

詳しくは省略しますが、ここまでわかれば、線形漸化式のときと同じように高速きたまさ法で $O(d\log d\log n)$ で元の問題を解くことができます。

ステップ3. 固有多項式で割ったあまりを求める

きたまさ法を使った解法では $f(x)$ を求めるために $O(d\log d\log n)$ の時間がかかりますが、実はこの問題の場合は $f(x)$ の係数を直接計算することができるので、よりよい計算量で解くことができます。
ここから先はひたすら計算が続きます。ちゃんとついてきてくださいね。

$f$ の形から、定数 $c$ と、 $d$ 次以下の多項式 $g$ を用いて、 $f(x)=c(x-r)^{d+1}+g(x)$ と書けます。この $c$ と $g$ を頑張って求めます。

$g$ について

この $g$ は要するに $(f\bmod (x-r)^{d+1})=(x^n\bmod (x-r)^{d+1})$ です。 $n< d+1$ なら自明です。そうでないときは $g(x+r)=( (x+r)^n\bmod x^{d+1})=\sum_{k=0}^{d}\binom{n}{k}r^{n-k}x^k$ になるので、あとは $x$ をずらすだけですね。
$\begin{eqnarray} g(x)&=&\sum_{k=0}^{d}\binom{n}{k}r^{n-k}(x-r)^k\\ &=&\sum_{k=0}^{d}\binom{n}{k}r^{n-k}\sum_{i=0}^{k}\binom{k}{i}(-r)^{k-i}x^i\\ &=&\sum_{i=0}^{d}r^{n-i}x^i(-1)^i\sum_{k=i}^{d}\binom{n}{k}\binom{k}{i}(-1)^k \end{eqnarray}$
$k$ に関する和が残ってしまいました。二項係数なので階乗に書き下して変形してもよいですが、組合せ的な言い換えをした方がわかりやすいことも多いです。今回、二項係数の積の部分は「 $n$ 個から $k$ 個選び、さらに選んだ $k$ 個から $i$ 個選ぶ方法の数」だと思えます。これは「 $n$ 個から $i$ 個選び、さらに選ばなかった $n-i$ 個から $k-i$ 個選ぶ方法の数」と等しいので、こういうふうに変形できますね。
$\begin{eqnarray} & &\sum_{k=i}^{d}\binom{n}{k}\binom{k}{i}(-1)^k\\ &=&\sum_{k=i}^{d}\binom{n}{i}\binom{n-i}{k-i}(-1)^k\\ &=&\binom{n}{i}(-1)^i\sum_{k=0}^{d-i}\binom{n-i}{k}(-1)^k\\ &=&\binom{n}{i}(-1)^i\binom{n-i-1}{d-i}(-1)^{d-i} \end{eqnarray}$
最後の等号は、 $\binom{x}{y}=\binom{x-1}{y-1}+\binom{x-1}{y}$ を使って $\binom{n-1}{*}$ を全部 $\binom{n-1-i}{*}$ に置き換えて計算するとわかります。そういうわけで、 $g$ は具体的に次のような形で書けました。
$g(x)=\sum_{i=0}^{d}r^{n-i}(-1)^{d-i}\binom{n}{i}\binom{n-i-1}{d-i}x^i$

$c$ について

すでに $g$ が求まっているので、適当な値を $x$ に代入して、恒等式 $f(x)=c(x-r)^{d+1}+g(x)$ を $c$ に関する $1$ 次方程式だと思って解くだけですが、適当な $x$ に対して $f(x)$ を求めるのはそんなに簡単ではないので、簡単に計算できる $x$ を探す必要があります。どんな $x$ なら簡単に求められるでしょうか。実は、 $(x-1)|p(x)$ と $f(x)=(x^n\bmod p(x))$ をあわせると $f(1)=(f(x)\bmod (x-1))=(x^n \bmod (x-1))=1$ になるので、 $x=1$ を選べばよいです。このとき、 $1=c(1-r)^{d+1}+g(1)$ を得ます。 $r\neq 1$ ならこれで話は終わりです。 $r=1$ のときは $c$ の係数が $0$ になってしまうので、これではうまくいきません。 $r=1$ のとき、元の問題は単なる単項式の和なので、最初から場合分けをして取り除いておいてもいいんですが、せっかくなのでいまやったのと同じような方法で求められないか考えてみましょう。

$r=1$ のときに $c$ を求める

$f(x)=c(x-r)^{d+1}+g(x)$ という恒等式を思い出すと、 $g$ は $d$ 次以下だったことから、 $c$ は $f$ の $x^{d+1}$ 次の係数と等しくなります。 $r=1$ のとき $p(x)=(x-1)^{d+2}$ なので、「 $f(x)=(x^n \mod (x-1)^{d+2})$ の $d+1$ 次の係数を求めよ」という問題になりますが、上で $g$ を求めるときにやったのと全く同じ議論で $\binom{n}{d+1}$ になることがすぐわかります。

ステップ4. $a_n$ を求める

これも線形漸化式の第 $n$ 項を求めるときと全く同じですが、一応説明しておきましょう。
求めたいものは $A_n=f(M)A_0$ の最後の成分でした。 $f(x)=\sum_{i=0}^{d+1}c_ix^i$ とおくと、 $A_n=\sum_{i=0}^{d+1}c_iM^iA_0=\sum_{i=0}^{d+1}c_iA_i$ となります。つまり $a_n=\sum_{i=0}^{d+1}c_ia_i$ なので、 $a_0,\ldots,a_{d+1}$ を予め愚直に求めておけば計算できます。

計算量解析

計算する必要があるものは $a_0,\ldots,a_{d+1}$ と $f(x)=c(x-r)^{d+1}+g(x)=\sum_{i=0}^{d+1}c_ix^i$ です。
計算量は $O(d \log {\rm MOD}+\log n)$ に見えますが、線形篩・線形逆元列挙・N乗数の高速列挙・二項係数の差分計算など、細かい工夫をたくさんすると $O(d+d \log {\rm MOD}/\log d+\log n)$ になります。