はじめに

この記事は 2022年夏のブログリレー 21 日目の記事です。

こんにちは。20B の tatyam です。

今日は数え上げ問題を解く上で強力な武器となる、母関数を用いた考察テクニックを習得していきましょう。

母関数を用いた考察テクニックは maspy さんの記事 以来広く知られるようになりました。(たぶん)

traP 内の各班では、有志が講師となって 講習会 を行っています。
以下は、最近私が行った FPS 講習会の資料ほぼそのままとなっています。

FPS 講習会

今日は aising2020 F - Two Snuke (Diff : 2741) を解けるようになっていきましょう。

母関数 (generating function)

母関数とは？

数列 $A=(A_0,A_1,A_2,\dots )$ の 母関数 を

$$A(x)=A_0+A_{1}x+A_2{x}^2+\cdots =\sum _{n=0}^{\infty }{A}_n{x}^n$$

と定義する。

数列について注意

数列は非負整数に対して数が $1$ つ定まっている (長さ加算無限の) 数列を考えることにする。
有限の数列を考えたい場合は残りを $0$ で埋めて無限数列であるものとする。

母関数の例

• 数列 $A=(3,2,1,0,0,0,\dots )$ の母関数は $A(x)=x^2+2x+3$ である。
• 数列 $B=(1,1,1,\dots )$ の母関数は $B(x)=1+x+x^2+\cdots$ である。
• 数列 $C=(\frac{1}{0!},\frac{1}{1!},\frac{1}{2!},\dots )$ の母関数は $C(x)=\frac{1}{0!}+\frac{1}{1!}x+\frac{1}{2!}{x}^2+\cdots =e^x$ である。

母関数は一般に無限個の項があるので、多項式というより 冪級数 (power series) である。(実際、$e^x$ は多項式ではない。)

この母関数が、数え上げの問題を解くのにとても役に立つ。

演習

1. 母関数が $\frac{x^2}{1-x}$ である数列の先頭 $5$ 項を答えよ。
2. 母関数が $\frac{1}{1-x^2}$ である数列の先頭 $5$ 項を答えよ。
3. 母関数が $\cos (x)$ である数列の先頭 $5$ 項を答えよ。

まとめ

• 数列に対して母関数 $A(x)=\sum _{n=0}^{\infty }{A}_n{x}^n$ を考える。
• 母関数をマクローリン展開すると数列が復元できる。
• 通常の関数と同じように変形してしまっても対応する数列が変わらないので OK

母関数と線形漸化式

線形漸化式から母関数

問題

数列 $\{\begin{array}{l}{F}_0=1\\ F_1=1\\ F_n=F_{n-1}+F_{n-2}\end{array}$ の母関数を求めよ。

解答

$F_n=F_{n-1}+F_{n-2}$ なので、$F(x)$ に $x$ を掛けて $1$ 要素ずらしたものと、$x^2$ を掛けて $2$ 要素ずらしたものを考える。

$$\begin{array}{rlr}F(x)& =& 1+1x+2x^2+3x^3+5x^4+8x^5+\cdots \\ xF(x)& =& 1x+1x^2+2x^3+3x^4+5x^5+\cdots \\ x^{2}F(x)& =& 1x^2+1x^3+2x^4+3x^5+\cdots \\ (1-x-x^2)F(x)& =& 1\end{array}$$

より、$F(x)=\frac{1}{1-x-x^2}$ である。

母関数から線形漸化式

問題

母関数が $\frac{1}{1-x-x^2}$ である数列の先頭 $5$ 項を求めよ。

解答

先頭の項から決定していこう。

$$A(x)=\sum _{n=0}^{\infty }{F}_n{x}^n=\frac{1}{1-x-x^2}$$

と置くと、

$$\left(F_0+F_{1}x+F_2{x}^2+F_3{x}^3+\cdots \right)(1-x-x^2)=1$$

$$\begin{array}{rlr}F(x)& =& F_0+F_{1}x+F_2{x}^2+F_3{x}^3+F_4{x}^4+F_5{x}^5+\cdots \\ -xF(x)& =& -F_{0}x-F_1{x}^2-F_2{x}^3-F_3{x}^4-F_4{x}^5-\cdots \\ -x^{2}F(x)& =& -F_0{x}^2-F_1{x}^3-F_2{x}^4-F_3{x}^5-\cdots \\ (1-x-x^2)F(x)& =& 1\end{array}$$

$0$ 次の項をみると $F_0=1$ なので、 $F_0=1$ に決定

$$\begin{array}{rlr}F(x)& =& 1+F_{1}x+F_2{x}^2+F_3{x}^3+F_4{x}^4+F_5{x}^5+\cdots \\ -xF(x)& =& -1x-F_1{x}^2-F_2{x}^3-F_3{x}^4-F_4{x}^5-\cdots \\ -x^{2}F(x)& =& -1x^2-F_1{x}^3-F_2{x}^4-F_3{x}^5-\cdots \\ (1-x-x^2)F(x)& =& 1\end{array}$$

$1$ 次の項をみると $F_1-1=0$ なので、 $F_1=1$ に決定

$$\begin{array}{rlr}F(x)& =& 1+1x+F_2{x}^2+F_3{x}^3+F_4{x}^4+F_5{x}^5+\cdots \\ -xF(x)& =& -1x-1x^2-F_2{x}^3-F_3{x}^4-F_4{x}^5-\cdots \\ -x^{2}F(x)& =& -1x^2-1x^3-F_2{x}^4-F_3{x}^5-\cdots \\ (1-x-x^2)F(x)& =& 1\end{array}$$

$2$ 次の項をみると $F_2-1-1=0$ なので、 $F_2=2$ に決定

つまり…？

$F_0=1$

$F_1-F_0=0$ より、$F_1=1$ に決定

$F_2-F_1-F_0=0$ より、$F_2=2$ に決定

$F_3-F_2-F_1=0$ より、$F_3=3$ に決定

$F_4-F_3-F_2=0$ より、$F_4=5$ に決定

つまり…？

が対応している。

母関数から線形漸化式の例

問題

母関数が $\frac{1-x^2}{1-2x}$ である数列の先頭 $5$ 項を計算せよ。

解答 $1$

$$\begin{array}{rl}\frac{1}{1-2x}& =1+2x+4x^2+8x^3+16x^4+\cdots \\ \frac{-x^2}{1-2x}& =0+0x-1x^2-2x^3-4x^4+\cdots \\ \frac{1-x^2}{1-2x}& =1+2x+3x^2+6x^3+12x^4+\cdots \end{array}$$

解答 $2$

$1-2x$ で割る $⟺$ もらう DP A[n] += A[n - 1] * 2

を利用する。

注 : 配る DP / もらう DP は計算順序の違いであって DP としては等価なのでどちらでもよい

問題

母関数が $\frac{1}{1-2x+x^2}$ である数列の先頭 $5$ 項を計算せよ。

解答 $1$

$1-2x+x^2$ で割る $⟺$ もらう DP A[n] += A[n - 1] * 2 - A[n - 2]

を利用する。

解答 $2$

$\frac{1}{1-2x+x^2}=\frac{1}{(1-x{)}^2}$

$\frac{1}{1-x}$ に対応する数列は $(1,1,1,1,1,\dots )$ なので、これを $1-x$ で割る $⟺$ もらう DP A[n] += A[n - 1] $⟺$ 累積和

$(1,1,1,1,1,\dots )$ の累積和を取って、$(1,2,3,4,5,\dots )$ が答え。

演習

4. 母関数が $\frac{1-x^4}{1-x}$ である数列の先頭 $5$ 項を計算せよ。
5. 母関数が $\frac{2}{2-x}$ である数列の先頭 $5$ 項を計算せよ。
6. 数列 $\{\begin{array}{l}{T}_0=1\\ T_1=1\\ T_2=1\\ T_n=T_{n-1}+T_{n-2}+T_{n-3}\end{array}$ の母関数を求めよ。

まとめ

• 母関数を $1-ax-b{x}^2+\cdots$ で割ることは、もらう DP
  A[n] += A[n - 1] * a + A[n - 2] * b + … を計算することに対応する。
• 特に、母関数を $1-x$ で割ることは、累積和に対応する。

数列を母関数に

問題

数列 $A=(1,1,2,2,3,3,4,4,\dots )$ の母関数を求めよ。

解答

同じ数字が $2$ 個連続しているから、数列を $2$ つに分ける :

$$\begin{gathered} B=(1,0,2,0,3,0,4,0,\dots ) \\ C=(0,1,0,2,0,3,0,4,\dots ) \end{gathered}$$

と、$C(x)=xB(x),A(x)=(1+x)B(x)$ が成り立つ。さらに、数列を詰める :

$$D=(1,2,3,4,\dots )$$

と、$D(x^2)=B(x)$ が成り立つ。
$D(x)=\frac{1}{(1-x{)}^2}$ より、求める母関数は $\frac{1+x}{(1-x^2{)}^2}$ である。

問題

数列 $A=(1,3,7,15,31,63,\dots )$ の母関数を求めよ。

解答

$$B=(1,1,1,1,1,1,\dots )$$

を足すと

$$A+B=(2,4,8,16,32,64,\dots )$$

となるから、$B(x)=\frac{1}{1-x},A(x)+B(x)=\frac{2}{1-2x}$ より、

$$A(x)=\frac{2}{1-2x}-\frac{1}{1-x}=\frac{1}{(1-x)(1-2x)}$$

演習

7. 数列 $A=(0,1,4,9,16,25,\dots )$ の母関数を求めよ。(ヒント : 階差を取ってみよう)

数え上げを母関数に

数列を母関数に変換することで最もうれしいことは、畳み込みが簡単に表現できることである。

畳み込みの例

問題

(互いに区別できない) $1$ 円硬貨を $4$ 枚、(互いに区別できない) $2$ 円硬貨を $2$ 枚持っているとき、$n$ 円を支払う方法は何通りあるか。

解答

$1$ 円硬貨のみを使って $x$ 円を支払う方法の数を $A=(1,1,1,1,1)$
$2$ 円硬貨のみを使って $y$ 円を支払う方法の数を $B=(1,0,1,0,1)$
とおくと、

$$C_n=\sum _{i+j=n}{A}_i{B}_j$$

が答えである。
これを実際に計算すると、$C=(1,1,2,2,3,2,2,1,1)$ となる。

畳み込み

数列 $A,B$ から新たな数列

$$C_n=\sum _{i+j=n}{A}_i{B}_j$$

を得る操作を 畳み込み (convolution) という。

畳み込みは冪級数の掛け算と対応していて、それぞれの母関数 $A(x),B(x),C(x)$ に対して $C(x)=A(x)\times B(x)$ が成り立つ。

有限数列の場合を考えれば多項式の話になってわかりやすい。
例えば、$A=(1,2,3),B=(1,-1,1)$ の場合を考えよう。
$A(x)\times B(x)=(1+2x+3x^2)\times (1-x+x^2)$ を展開する。

例えば、$A(x)\times B(x)$ の $x^2$ の係数を求めるには、斜めに並んだ $3x^2,-2x^2,x^2$ を足して $2x^2$ とすればよい。
この計算は、畳み込みの式 $C_2=\sum _{i+j=2}{A}_i{B}_j$ にほかならない。

畳み込みを使って数える

問題

非負整数 $1$ つが出るルーレットがあり、ルーレットを回すと確率 $P_n$ で $n$ が出る。
このルーレットを $10$ 回回したときの出た数の合計が $m$ である確率を $Q_m$ とするとき、母関数 $Q(x)$ を母関数 $P(x)$ を使って表せ。

解答

$2$ 回回したときに和が $n$ になる確率は $\sum _{i+j=n}{P}_i{P}_j$ で、この母関数は $(P(x){)}^2$ である。
$P^{\prime }(x)=(P(x){)}^2$ と置くと、$3$ 回回したときに和が $n$ になる確率は、$i+j=n$ なる $(i,j)$ についての、前 $2$ 回の和が $i$ で後ろ $1$ 回が $j$ になる確率の和であるから、$\sum _{i+j=n}{P}_i^{\prime }{P}_j$ で、この母関数は $(P(x){)}^3$ である。
これを繰り返せば、$Q(x)=(P(x){)}^{10}$ がわかる。

演習

8. (互いに区別できない) $1$ 円硬貨を無限枚、(互いに区別できない) $2$ 円硬貨を無限枚持っているとき、$n$ 円を支払う方法の数を $A_n$ とする。母関数 $A(x)$ を求めよ。

問題を母関数に帰着する

形式的冪級数 $f(x)$ に対して、$[x^n]f(x)$ で $f(x)=\sum _{i=0}^{\infty }{f}_i{x}^i$ と展開したときの $x^n$ の係数 $(f_n)$ を表す。

問題

https://atcoder.jp/contests/dp/tasks/dp_m

$N$ 人の子供に $K$ 個の区別できない飴を配る。子供 $i$ には $0$ 個以上 $A_i$ 個以下の飴を配らなければならない。配り方は何通りか？

解答

飴の 合計が $K$ 個という和の制約があるので、飴の個数が $n$ 個の場合の配り方が $x^n$ の係数となる母関数を考える。

子供 $1$ へ飴を $n$ 個配る方法の数の母関数は、$\sum _{n=0}^{A_1}{x}^n=\frac{1-x^{A_1+1}}{1-x}$
子供 $2$ へ飴を $n$ 個配る方法の数の母関数は、$\sum _{n=0}^{A_2}{x}^n=\frac{1-x^{A_2+1}}{1-x}$
子供 $1,2$ へ飴を合計 $n$ 個配る方法の数の母関数は、畳み込みになって $\frac{1-x^{A_1+1}}{1-x}\cdot \frac{1-x^{A_2+1}}{1-x}$
子供 $1,2,3$ へ飴を合計 $n$ 個配る方法の数の母関数は、畳み込みになって $\frac{1-x^{A_1+1}}{1-x}\cdot \frac{1-x^{A_2+1}}{1-x}\cdot \frac{1-x^{A_3+1}}{1-x}$

$N$ 人の子供へ飴を合計 $n$ 個配る方法の数の母関数は、$\prod _{i=1}^N\frac{1-x^{A_i+1}}{1-x}$ なので、$[x^K]\left(\prod _{i=1}^N\frac{1-x^{A_i+1}}{1-x}\right)$ が答えである。

問題

長さ $N$ の数列 $A=(A_1,\dots ,A_N)$ がある。
数列 $X=(X_1,\dots ,X_{|X|})$ に対して、その うれしさ を総積 $\prod _{i=1}^{|X|}{X}_i$ で定義する。
数列 $A$ の部分列のうち長さが $M$ であるもの全てに対してうれしさを計算し、その総和を求めよ。

部分列 : 数列から要素をいくつか ($0$ 個以上) 削除したもの

解答

$A$ の各要素に対して 使う・使わない を選び、ちょうど $M$ 個 の要素を使うときの総積という形になっている。

そこで、ちょうど $m$ 個使うような選び方全てに対してのうれしさの和が $x^m$ の係数となる母関数を考える。
$N=1$ の場合を考えると、母関数は $1+A_{1}x$ であるから、これを畳み込んで

$$[x^M]\left(\prod _{i=1}^{N}1+A_{i}x\right)$$

が答えである。

実際、$N=2$ のとき、

$$1+(A_1+A_2)x+A_1{A}_2{x}^2=(1+A_{1}x)(1+A_{2}x)$$

$N=3$ のとき、

$$1+(A_1+A_2+A_3)x+(A_1{A}_2+A_1{A}_3+A_2{A}_3)x^2+A_1{A}_2{A}_3{x}^3=(1+A_{1}x)(1+A_{2}x)(1+A_{3}x)$$

となるから、成り立っていることがわかる。

問題

正整数 $N,M,K$ が与えられる。以下の条件を満たす正整数列 $A=(A_1,A_2,\dots ,A_{|A|})$ はいくつ存在するか？

• $1\le |A|\le N$
• $1\le A_i\le M(1\le i\le |A|)$
• $\left(\sum _n{A}_n\right)\le K$

解答

総和が $K$ 以下 という和の制約があるので、総和が $n$ となる数列の数が $x^n$ の係数となるような母関数を考える。

$|A|=1$ の場合、母関数は $\sum _{n=1}^M{x}^n=\frac{x-x^{M+1}}{1-x}$

これを $f(x)$ とすると、$|A|=n$ の場合の母関数は $(f(x){)}^n$

実際は $1\le |A|\le N$ であるから総和を取って、目的の母関数は

$$\begin{array}{rl}\sum _{n=1}^N(f(x){)}^n& =\frac{f(x)-(f(x){)}^{N+1}}{1-f(x)}\\ & =\frac{x-x^{M+1}}{(1-x)-(x-x^{M+1})}-\frac{(x-x^{M+1})(f(x){)}^N}{(1-x)-(x-x^{M+1})}\\ & =\frac{x-x^{M+1}}{1-2x+x^{M+1}}-\frac{(x-x^{M+1}{)}^{N+1}}{(1-x{)}^N(1-2x+x^{M+1})}\end{array}$$

したがって、

$$\sum _{k=1}^K[x^k]\left(\frac{x-x^{M+1}}{1-2x+x^{M+1}}-\frac{(x-x^{M+1}{)}^{N+1}}{(1-x{)}^N(1-2x+x^{M+1})}\right)$$

が答え。$\sum$ が残っていると嫌なので、$1-x$ で割る $⟺$ 累積和を取って、

$$[x^K]\left(\frac{x-x^{M+1}}{(1-x)(1-2x+x^{M+1})}-\frac{(x-x^{M+1}{)}^{N+1}}{(1-x{)}^{N+1}(1-2x+x^{M+1})}\right)$$

とする。

ちなみに : $\frac{f(x)-(f(x){)}^{N+1}}{1-f(x)}$ の時点で $O(K\log K)$ 時間で求まりそうなことがわかって、最後の式から $O(K)$ 時間で求まりそうなことがわかる。

まとめ

• 「総和がちょうど $X$」や「総和が $X$ 以下」のような総和の制約があり、確率 や 組み合わせの数 など、掛け算をするものを数える場合、母関数の掛け算で表すことができる。
• 「どのように DP をするかを考察する」という難しいパートが「母関数に変換する」→「簡単な式に変形する」→「式から計算方法を復元する」となり、比較的機械的にできるようになる。

母関数から計算方法を復元する

演習

9. $(N,K)$ の小さいところで $[x^K]\frac{1}{(1-x{)}^N}(N\ge 1)$ の値を計算し、一般項を推測せよ。

二項係数

注 : 説明の都合上、数列とその母関数を同一視する。

$\frac{1}{(1-x{)}^n}$ の前 $N$ 項は $O(N)$ 時間で求められる。

$1-x$ で割る $⟺$ 累積和であるから、

これはパスカルの三角形になっているから、

$\frac{1}{(1-x{)}^n}=\sum _{k=0}^{\infty }\left(\genfrac{}{}{0ex}{}{n-1+k}{k}\right)x^k$

掛け算 (sparse)

数列 $A(x)=\sum _{n=0}^{\infty }{A}_n{x}^n$ に $K$ 個の項以外が全て $0$ であるような数列 $B(x)=\sum _{k=1}^K{b}_k{x}^{n_k}$ を掛けたときの前 $N$ 項は $O(NK)$ 時間で求められる。

普通に展開すればよい。

例

$A(x)=4-3x+2x^2-x^3-x^4-4x^5+5x^6$ に $B(x)=2+x^2$ を掛けた数列を求めよ。

割り算 (sparse)

数列 $A(x)=\sum _{n=0}^{\infty }{A}_n{x}^n$ を $K$ 個の項以外が全て $0$ であるような数列 $B(x)=\sum _{k=1}^K{b}_k{x}^{n_k}$ で割ったときの前 $N$ 項は $O(NK)$ 時間で求められる。

分母・分子に適当なものを掛けて分母を $1-f(x)$ の形にして、DP に変換すればよい。

例

$A(x)=4-3x+2x^2-x^3-x^4-4x^5+5x^6$ を $B(x)=1+2x^2$ を割った数列の前 $8$ 項を求めよ。

掛け算

数列 $A(x)=\sum _{n=0}^{\infty }{A}_n{x}^n$ に数列 $B(x)=\sum _{n=0}^{\infty }{B}_n{x}^n$ を掛けたときの前 $N$ 項は $O(N\log N)$ 時間で求められる。

FFT をすればよい : 【競プロer向け】FFT を習得しよう！ | 東京工業大学デジタル創作同好会traP

AC Library にも入っている : Convolution | AC Library

inv, log, exp of FPS

数列 $A(x)=\sum _{n=0}^{\infty }{A}_n{x}^n$ に対して、

• $\frac{1}{A(x)}$ の前 $N$ 項は $O(N\log N)$ 時間で求められる。
• $\log (A(x))=-\sum _{k=1}^{\infty }\frac{(1-A(x){)}^k}{k}=\int \frac{A^{\prime }(x)}{A(x)}$ の前 $N$ 項は $O(N\log N)$ 時間で求められる。
• $\exp (A(x))=\sum _{k=0}^{\infty }\frac{(A(x){)}^k}{k!}$ の前 $N$ 項は $O(N\log N)$ 時間で求められる。

詳しくは 形式的冪級数（FPS）の inv，log，exp，pow の定数倍の軽いアルゴリズム | opt の競プロブログ などを参照

pow of FPS

数列 $A(x)=\sum _{n=0}^{\infty }{A}_n{x}^n$ と数 $K$ に対して、

• $(A(x){)}^K=\exp (K\log (A(x)))$ の前 $N$ 項は $O(N\log N)$ 時間で求められる。

これは $\log$ と $\exp$ ができればよい。

演習

10. 長さ $N$ の数列 $A(x)=\sum _{n=0}^{N-1}{A}_n{x}^n$ と長さ $M$ の数列 $B(x)=\sum _{n=0}^{M-1}{B}_n{x}^n$ を受け取り、$\frac{A(x)}{B(x)}$ の前 $N$ 項を $O(NM)$ 時間で計算するコードを書け。$B_0=1$ であると仮定して良い。

入力例1

7 3
4 -3 2 -1 -1 -4 5
1 0 2

出力例1

4 -3 -6 5 11 -14 -17

入力例2

8 4
1 0 0 0 0 0 0 0
1 -1 -1 -1

出力例2

1 1 2 4 7 13 24 44

入力例3

8 4
1 7 21 35 35 21 7 1
1 3 3 1

出力例3

1 4 6 4 1 0 0 0

母関数から計算方法を復元する例

問題

長さ $N+1$ の数列 $A(x)=\sum _{n=0}^N{A}_n{x}^n$ と整数 $K\ge 0$ が与えられる。$[x^N]\frac{A(x)}{(1-x{)}^K}$ を $O(N)$ 時間で求めよ。

注 : 数がどれだけ大きくても四則演算は $O(1)$ 時間でできるものとする (よくあるのは $\mathrm{mod}998244353$ で計算する)

解答

$\frac{1}{(1-x{)}^K}$ の前 $N+1$ 項は $O(N)$ 時間で計算できる。
畳み込みで $A(x)\times \frac{1}{(1-x{)}^K}$ の前 $N+1$ 項を求めるには $O(N\log N)$ 時間かかるが、$N$ 項目だけなら $O(N)$ 時間で計算できる。

実装

N, K = map(int, input().split())
A = list(map(int, input().split()))

# 1 / (1 - x)^K の前 N + 1 項を O(N) 時間で求める
# B[n] = [x^n](1 - x)^-K = binom(n - 1 + K, n - 1)
B = [1]
for i in range(N):
    B.append(B[i] * (N - 1 + K - i) // (i + 1))

# [x^N] (A * B) を O(N) 時間で求める
ans = 0
for i in range(N + 1):
    ans += A[i] * B[N - i]

print(ans)

問題

https://atcoder.jp/contests/dp/tasks/dp_m

$N$ 人の子供に $K$ 個の区別できない飴を配る。子供 $i$ には $0$ 個以上 $A_i$ 個以下の飴を配らなければならない。配り方は何通りか？

$[x^K]\left(\prod _{i=1}^N\frac{1-x^{A_i+1}}{1-x}\right)$ を $O(NK)$ 時間で求めよ。

解答

必要なのは $K$ 項目であるから、計算途中において $K$ 項目より先を計算する必要はない。

$(1-x^{A_i+1})$ を掛ける操作、$(1-x)$ で割る操作は $O(K)$ 時間でできるから、全体で $O(NK)$ 時間である。

ポイント : 本来、「累積和を使って DP する」という考察が必要なところが機械的にできている。

問題

$[x^K]\left(\prod _{i=1}^N\frac{1-x^{A_i+1}}{1-x}\right)$ を $O(N{2}^N)$ 時間で求めよ。

解答

今度は $(1-x^{A_i+1})$ を掛ける操作、$(1-x)$ で割る操作に $O(K)$ 時間かけることができない。

$$\prod _{i=1}^N\frac{1-x^{A_i+1}}{1-x}=\frac{{\prod }_{i=1}^N(1-x^{A_i+1})}{(1-x{)}^N}$$

と変形すると、$\frac{1}{(1-x{)}^N}$ は二項係数になるから簡単に計算できる。
分子の $\prod _{i=1}^N(1-x^{A_i+1})$ の前 $K+1$ 項は $O(NK)$ 時間で計算できるが、$\prod _{i=1}^N(1-x^{A_i+1})$ を展開すると高々 $2^N$ 項しかないことから、$K$ に関係なく $O(2^N)$ 時間で計算できる。

分子が $\prod _{i=1}^N(1-x^{A_i+1})=\sum _{i=0}^{2^N-1}{B}_i{x}^{n_i}$ と展開されるものとすると、

$$\begin{array}{rl}[x^K]\frac{{\sum }_{i=0}^{2^N-1}{B}_i{x}^{n_i}}{(1-x{)}^N}& =\sum _{i=0}^{2^N-1}\left(B_i\times [x^{K-n_i}]\frac{1}{(1-x{)}^N}\right)\\ & =\sum _{i=0}^{2^N-1}{B}_i\left(\genfrac{}{}{0ex}{}{N-1+K-n_i}{K-n_i}\right)\\ & =\sum _{i=0}^{2^N-1}{B}_i\left(\genfrac{}{}{0ex}{}{N-1+K-n_i}{N-1}\right)\end{array}$$

$\left(\genfrac{}{}{0ex}{}{N-1+K-n_i}{N-1}\right)$ は $O(N)$ 時間で求められるから、全体で $O(N{2}^N)$ 時間である。

ポイント : 本来、包除原理を使った考察が必要なところが機械的にできている。

問題

数列 $A(x)=\sum _{n=0}^{\infty }{A}_n{x}^n$ に対して、$\sum _{n=0}^{K-1}(A(x){)}^n=\frac{1-(A(x){)}^K}{1-A(x)}$ の前 $N$ 項を $O(N\log N)$ 時間で求めよ。

解答

$(A(x){)}^K$ の前 $N$ 項は pow of FPS を使って $O(N\log N)$ 時間で求まる。
$\frac{1}{1-A(x)}$ の前 $N$ 項は inv of FPS を使って $O(N\log N)$ 時間で求まる。
これを畳み込めば、$\frac{1-(A(x){)}^K}{1-A(x)}$ の前 $N$ 項は $O(N\log N)$ 時間で求まる。

$[x^K]\frac{P(x)}{Q(x)}$ を求める

行列累乗

長さ $N$ の数列 $P(x)=\sum _{n=0}^{N-1}{P}_n{x}^n$ と長さ $N+1$ の数列 $Q(x)=\sum _{n=0}^N{Q}_n{x}^n$ に対して、$[x^K]\frac{P(x)}{Q(x)}$ は $O(N^3\log K)$ 時間で求められる。

前 $N$ 項を求めておき、分母を線形漸化式に変換して、$1$ 要素シフトする行列を作り、行列累乗をすればよい。

例

$[x^K]\frac{1}{1-ax-b{x}^2}$ を求める。
分母が $1-ax-b{x}^2$ $⟺$ 線形漸化式 $A_n=a{A}_{n-1}+b{A}_{n-2}$ であるから、

$$\left[\begin{array}{cc}{A}_{n-2}& A_{n-1}\end{array}\right]\left[\begin{array}{cc}0& b\\ 1& a\end{array}\right]=\left[\begin{array}{cc}{A}_{n-1}& A_n\end{array}\right]$$

が成り立ち、

$$\left[\begin{array}{cc}1& a\end{array}\right]{\left[\begin{array}{cc}0& b\\ 1& a\end{array}\right]}^n=\left[\begin{array}{cc}{A}_n& A_{n+1}\end{array}\right]$$

である。これは行列累乗で $O(\log K)$ 時間で計算できる。

高速きたまさ法 (Feduccia's algorithm)、$Q(x)Q(-x)$ (Bostan–Mori algorithm)

↑ は $O(N\log N\log K)$ 時間で計算できる。

詳しくは 線形漸化式の高速計算 | Nyaan’s Library などを参照

演習

今なら aising2020 F - Two Snuke が解けるはず…！

11. $0\le s_1<s_2,s_1+s_2=n$ を満たす整数の組 $(s_1,s_2)$ 全てについて $(s_2-s_1)$ を合計したものを $F_n$ とする。$(F_0,F_1,\dots ,F_6)$ を求めよ。
12. 母関数 $F(x)$ を求めよ。(ヒント : 階差を取ってみよう)
13. 元の問題 の答えを $[x^N]f(x)$ の形式で表せ。
14. どれくらいの時間計算量で解けるか？

おまけ

オンライン整数列大辞典

有名な整数列がたくさん載っている辞典。母関数がわからないとき、一般項がわからないとき、ここで調べるとみつかるかも？

https://oeis.org/search?q=0%2C1%2C2%2C4%2C6%2C9%2C12

b(n) = A002620(n+2) = number of multigraphs with loops on 2 nodes with n edges [so g.f. for b(n) is 1/((1-x)^2*(1-x^2))].

g.f. (Generating Functiion) で母関数が書いてあったり、漸化式が書いてあったり、一般項が書いてあったりする。

FFT マージテク

$M$ 個の多項式 $f_1,\dots ,f_M$ があり、次数の合計が $\sum _{i=1}^M\mathrm{deg}(f_i)=N$ であるとき、総積の係数 $\prod _{i=1}^M{f}_i(x)=\sum _{n=0}^N{A}_n{x}^n$ は $O(N\log N\log M)$ で求められる。

• $M$ 個の多項式のうち最も次数の小さなもの $2$ 個を掛けることを繰り返す (priority_queue で管理)
• $M$ 個の多項式のうち最も古いもの $2$ 個を掛けることを繰り返す (deque で管理)

などをするとできる。

実装

F = deque([f_1, …, f_M])
while len(F) > 1:
    a = F.popleft()
    b = F.popleft()
    F.append(convolution(a, b))

例

長さ $N$ の数列 $A=(A_1,\dots ,A_N)$ がある。
数列 $X=(X_1,\dots ,X_{|X|})$ に対して、その うれしさ を総積 $\prod _{i=1}^{|X|}{X}_i$ で定義する。
数列 $A$ の部分列のうち長さが $M$ であるもの全てに対してうれしさを計算し、その総和を求めよ。

$$[x^M]\left(\prod _{i=1}^{N}1+A_{i}x\right)$$

は $O(N(\log N{)}^2)$ で求められる。

exp-log

数列 $A(x)=\sum _{n=0}^{\infty }{A}_n{x}^n$ に対して、$\prod _{k=1}^{M}A(x^k)$ の前 $N$ 項は $O(N(\log N+\log M))$ 時間で求められる。

問題

$1$ から $N$ までの 分割数 を $O(N\log N)$ 時間で求めよ。

解答

分割数は以下の問題と同じである。

(互いに区別できない) $1$ 円硬貨を無限枚、$2$ 円硬貨を無限枚、$3$ 円硬貨を無限枚、$4$ 円硬貨を無限枚、……
持っているとき、$n$ 円を支払う方法は何通りか？

$m$ 円硬貨のみを用いて $n$ 円を支払う方法の数の母関数は $\frac{1}{1-x^m}$
$N$ 円より大きな硬貨には意味がないから、分割数の母関数は $\prod _{m=1}^N\frac{1}{1-x^m}$ である。