概要
AtCoder 第一回 日本最強プログラマー学生選手権 -予選- → ■
問題F – Candy Retribution → ■
自分の提出 → ■
私はこの問題は、コンテスト中には解ききれなかった（解き方は分かったものの、残り時間がなくて実装を終えられなかった）のですが、公式解説にはない、形式的べき級数による考察をとったので、考察の要点などを書き残しておこうと思います。また、この問題をきっかけとして、形式的べき級数の数え上げへの基本的な利用について、少し一般的に扱いました。
以下、二項係数 ${}_n\mathrm{C}_k$ を $\binom{n}{k}$ と表記します。
形式的べき級数・代表的な等式
多項式を無限に続けていった式
\[ f(T) = a_0 + a_1T + a_2T^2 + a_3T^3 + \cdots \] を形式的べき級数と言います。多項式の仲間で、和・差・積が定義できて、結合・交換・分配法則などが成り立ります。多項式の真似で定義しようとすると、個別の係数は有限和なので自然な定義が可能です。
細かいところを気にしない人はまぁ多項式だと思っておいていても困らないかもしれません。ただし、次数がいくらでも高い項がある（「最高次の項」に依存した議論が行えない）ことに起因して、いくつかの目立った違いが起こります。
・$f$ への代入 $f(t)$ は考えない、あるいは慎重に考える（無限級数）。
・$f$ の次数 $\deg(f)$ を用いる議論は通常行わない。そもそも通常 $\deg(f)$ を定義すらしない。
・多項式以外を $\deg(f) = \infty$ と定義することはできるが、$\deg(fg) = \deg(f) + \deg(g)$ など期待したい性質が成り立たない（以下の例）。
$f(T) = g(T)h(T)$ が成り立つとき、$g(T) = \frac{f(T)}{h(T)}$ と書きます。次のような等式が代表的です：
\[\begin{align*} \frac{1}{1-T} &= 1+T+T^2+T^3+T^4+\cdots = \sum_{n=0}^{\infty}T^n,\\
\frac{1}{(1-T)^2} &= 1+2T+3T^2+4T^3+5T^4+\cdots = \sum_{n=0}^{\infty}(n+1)T^{n},\\
\frac{1}{(1-T)^3} &= 1+3T+6T^2+10T^3+15T^4+\cdots = \sum_{n=0}^{\infty}\binom{n+2}{2}T^{n}.\end{align*}\]
より一般に、$\frac{1}{(1-T)^k} = \sum_{n=0}^{\infty}\binom{n+k-1}{k-1}T^n$ が成り立ちます。
例えばひとつ目の式は、分母のない形で書けば
\[ (1-T)(1+T+T^2+T^3+T^4+\cdots) = 1\] となります。左辺を分配法則でそれっぽく計算すると、$n$ 次の項（$n\geqq 1$）は $1\cdot T^n, (-T)\cdot T^{n-1}$ の2通りが出てきて、和が $0$ になりますね。
※ 等式は多項式の世界では、$n$ 次式と $m$ 次式の積が $n+m$ 次式になるため、このように定数じゃなさそうなものの積が定数となることはありません。形式的べき級数には最高次の項がないので、こういうことができます。
$\frac{1}{(1-T)^k}$ の場合の計算については後述します。
分割数と形式的べき級数
適当な整数の和として $n$ を作る方法を数え上げる問題は、分割数と呼ばれる問題に属します。
常識として、分割数といえば形式的べき級数です（公式 editorial の真似）。例えば次の数え上げを考えてみましょう。
【問題】$x+y+z = n$ となるような整数 $0\leqq x\leqq 2$, $0\leqq y\leqq 3$ および偶数 $0\leqq z\leqq 4$ の組 $(x,y,z)$ の個数 $a_n$ を求めよ。
【解答】$(1+T+T^2)(1+T+T^2+T^3)(1+T^2+T^4) = a_0+a_1T+a_2T^2+\cdots$ とするとき、$a_n$が求める数である。
1つめのカッコから $T^x$、2つめのカッコから $T^y$、3つめのカッコから $T^z$ を選ぶという要領です。このように、制約条件ごとにそれっぽい多項式、あるいは形式的べき級数を作って掛け合わせることにより、
「数え上げよ」 → 「形式的べき級数の積を計算せよ」 と問題を変形することができます。
「分割数といえば形式的べき級数」というのは誇張ではなく、数学研究においても、自然数への分割・平方数への分割・素数への分割などをこのように形式的べき級数に変換するというのは定跡です。
$x,y,z$ に余計な制約条件がない単純な場合には、次の形になります。
【問題】$x+y+z = n$ となるような非負整数 $x,y,z$ の組の個数 $a_n$ を求めよ。
【解答】$(1+T+T^2+T^3+T^4+T^5+\cdots)^3 = a_0+a_1T+a_2T^2+\cdots$ とするとき、$a_n$が求める数である。
$1+T+\cdots+T^n$ と $n$次で打ち切っても良いですが、$x,y,z$ として $n+1$ 以上も許すと左辺には多項式ではない形式的べき級数が登場します。そして、この形式的べき級数 $1+T+T^2+T^3+T^4+T^5+\cdots$ は、先述のように $\frac{1}{1-T}$ と簡潔に表せるのでした。したがって「$\frac{1}{(1-T)^3}$ を計算せよ」というのが課題です。
これを形式的べき級数の演算で計算することもできます（例えば微分するとか）が、今回は、この対応から形式的べき級数の等式を導きましょう。元の問題は高校数学標準な数え上げで、$\binom{n+2}{2}$ 通りが答になります（「重複組み合わせ」で検索すると見つかるかと）。したがって次の等式が得られます：
\[ \frac{1}{(1-T)^3} = \sum_{n=0}^{\infty}\binom{n+2}{2}T^n.\]
同じように考えると、先述の$\frac{1}{(1-T)^k}$の等式が得られます。
累積和と形式的べき級数
数列の累積和も、形式的べき級数との相性が良いです。無限数列 $A_0,A_1,A_2,\ldots$ に対して $S_n = A_0+A_1+\cdots+A_n$ と定義することにします。このとき次が成り立ちます：\[ (A_0+A_1T+A_2T^2+A_3T^3+\cdots)(1+T+T^2+T^3+\cdots) \\= S_0+S_1T+S_2T^2+S_3T^3+\cdots.\]左辺を展開して係数比較すれば確かめられます。特に $f_A(T) = \sum_{n=0}^{\infty}A_nT^n$, $f_S(T) = \sum_{n=0}^{\infty}S_nT^n$ とすると、$f_A(T)\cdot \frac{1}{1-T} = f_S(T)$ という簡潔な式で表せます（形式的べき級数 $f_A$, $f_S$ は数列 $A$, $S$ の母関数と呼ばれることもあります）。
「累積和をとる」というのが「$1-T$ で割る」という定数倍で表すことができました。
公式 editorial の常識も回収してみましょう。
【問題】$x+y+z \leqq n$ となるような非負整数 $x,y,z$ の組の個数 $S_n$ を求めよ。
【解答】$x+y+z=n$ となる $x,y,z$ の個数を $A_n$ とする。$A_n$ は $\frac{1}{(1-T)^3}$ の係数だった。したがって $S_n = A_0+\cdots + A_n$ は $\frac{1}{(1-T)^{3}}\cdot \frac{1}{1-T} = \frac{1}{(1-T)^4}$ の係数なので、$\binom{n+3}{3}$ に等しい。
もちろん形式的べき級数を使わない方法もあります。例えば $w = n – (x+y+z)$ と置くことで、問題の数え上げは、$x+y+z+w=n$ となる非負整数 $x,y,z,w$ の数え上げと対応がつくので典型的な重複組み合わせです。このような考察に基づく対応付けをせずとも、ただの有理式の恒等式という形で異なる数え上げを結び付けられるところが、面白いところでしょう。
第一回JSC-予選- 問題F – Candy Retribution
問題文 → ■
和 $A_1+\cdots+A_N$ の条件を忘れて、まずは数え上げる対象を整理します。余事象を考えます。上側集合、下側集合に入る番号を選ぶことで、次を数えて最後に $\binom{N}{M}$ 倍すればよいです。
(1)「$A_1,\ldots,A_M \quad >\quad A_{M+1},\ldots,A_N$ となる非負整数列のうちで、和が適切なものを数えよ」
数え上げでは、漏れ・重複なく数えるのが肝要です。そのためには明確な場合分けの基準を設けます。今回は、下側集合の最大値 $k$ で場合分けをします。$k$ ごとに次を数えます。
(2)「$A_1,\ldots,A_M > k, \quad\max\{A_{M+1},\ldots,A_N\}=k$ となる非負整数列のうちで、和が適切なものを数えよ」
さらに、最大値が $k$ という条件は、(すべて $k$ 以下) かつ not (すべて $k-1$ 以下) と書き換えられます。したがって、次の2つの数え上げの差分をとることにします。
(3) 「$A_1,\ldots,A_M > k,\quad k\geqq A_{M+1},\ldots,A_N$ となる非負整数列のうちで、和が適切なものを数えよ」
(4) 「$A_1,\ldots,A_M > k,\quad k-1\geqq A_{M+1},\ldots,A_N$ となる非負整数列のうちで、和が適切なものを数えよ」
(3) のような数列に対して、$n = A_1+\cdots+A_N$ となる非負整数列が $p_n$ 通りあるとしましょう。これを形式的べき級数に基づいて考察すると、次のようになります：
\[\begin{align*}
\sum_{n=0}^{\infty} p_nT^n &= (T^{k+1}+T^{k+2}+T^{k+3}+\cdots)^{M}(1+T+T^2+\cdots+T^k)^{N-M}\\
&=\biggl(\frac{T^{k+1}}{1-T}\biggr)^M\biggl(\frac{1-T^{k+1}}{1-T}\biggr)^{N-M}\\
&=\frac{1}{(1-T)^N}T^{(k+1)M}(1-T^{k+1})^{N-M}.
\end{align*}\]$A_1,\ldots,A_M$ については $k+1$ 以上から自由に選ぶという条件が、$T^{k+1}+T^{k+2}+T^{k+3}+\cdots$ に対応しています。$A_{M+1},\ldots,A_N$ については $k$ 以下から自由に選ぶという条件が、$1+T+T^2+\cdots+T^k$ に対応しています。(4)についても同様に考えて差分をとることで、(2) の数え上げの母関数は次のように表せます：
\[\frac{1}{(1-T)^N}T^{(k+1)M}\bigl((1-T^{k+1})^{N-M}-(1-T^k)^{N-M}\bigr).\](1) の数え上げの母関数は次のように表せます（$k$ について和をとるだけです）：
\[\frac{1}{(1-T)^N}\sum_{k=0}^{\infty}T^{(k+1)M}\bigl((1-T^{k+1})^{N-M}-(1-T^k)^{N-M}\bigr).\]さらに $\sum_{k=0}^{\infty}T^{(k+1)M}(1-T^{k+1})^{N-M} = \sum_{k=0}^{\infty} T^{kM}(1-T^k)^{N-M}$ を用いることで、(1) の数え上げの母関数は
\[\frac{1}{(1-T)^N}\sum_{k=0}^{\infty}(T^{kM}-T^{(k+1)M})(1-T^{k})^{N-M}
=\frac{1-T^M}{(1-T)^N}\sum_{k=1}^{\infty}T^{kM}(1-T^{k})^{N-M}\] となります。
記号や分かったことを整理しながら、解法の完成に向かいましょう。
・(1) のような数列で、$A_1+\cdots+A_N = n$ となるものの個数を $a_n$ と書く。知りたいのは $a_L+\cdots+a_R$ である。
・$a_n$ の累積和を $S_n = a_0+\cdots+a_n$ と書く。$S_R$ と $S_{L-1}$ が知りたい。
・$F(T) = T^M(1-T)^{N-M}$ とおく。
・$a_n$ の母関数は $\frac{1-T^M}{(1-T)^N}\sum_{k=1}^{\infty}T^{kM}(1-T^{k})^{N-M}=\frac{1-T^M}{(1-T)^N}\sum_{k=1}^{\infty}F(T^k)$ に等しい。
・したがって、$S_n$ の母関数は $\frac{1-T^M}{(1-T)^{N+1}}\sum_{k=1}^{\infty}F(T^k)$ に等しい。これの $T^R$ と $T^{L-1}$ の係数が知りたい。
・$P(T) = \sum_{k=1}^{\infty}F(T^k)$, $Q(T) = \frac{1-T^M}{(1-T)^{N+1}}$ と書く。$S_n$ の母関数は $P(T)Q(T)$ である。
あとは適切に計算するだけです。
まず、 $F(T)$ は二項定理により（階乗の事前計算のもと） $O(M)$ で計算できます。$P(T)$ は無限和で定義されていますが、これの $R$ 次以下の項だけが必要です。したがって有限和 $\sum_{k=1}^{R}$ でよいですし、各 $k$ に対して $F(T^k)$ が $R$ 次以下に寄与する項は $\lfloor\frac{R}{k}\rfloor$ 程度しかありません。よって $P(T) = \sum_{k=1}^{\infty}F(T^k)$ の計算は $O(R\log R)$ で行えます。
$Q(T)$ の計算は、まず $\frac{1}{(1-T)^{N+1}}$ は係数に二項係数 $\binom{n+N}{N}$ が並ぶのでした。よって階乗の事前計算のもと係数ごとに$O(1)$、$R$ 次以下の部分が $O(R)$ で計算できます。 $(1-T^{M})$ を乗ずるのも、$M$ 個ずれた部分を引くだけなので係数ごとに $O(1)$ です。
知りたいのは $P(T)Q(T)$ の個別の係数です。$R$ 次多項式の積は単純に $O(R^2)$ あるいはFFTなどを用いて $O(R\log R)$ ですが、今回は個別の係数を取り出すだけなので必要な組み合わせの積を全て足し合わせて $O(R)$ です。以上を実装すればACできます。ただしpythonだと適切にnumpyを使わないとTLEすると思います。
類題？
まだまだあると思いますが、すぐに見つかったものを挙げておきます。
◆Educational DP Contest：問題M – Candies
「$\frac{1-x^n}{1-x}$をかける」と思うか「累積和と差分でdp更新」と思うか。
◆ARC 028：問題D – 注文の多い高橋商店
dp更新と思うよりも、形式的べき級数の積と思う方が、交換法則や戻せることなどが明快？
◆square869120Contest #3：問題G – フィボナッチ数の総和
本記事では解説していませんが、線形漸化式と形式的べき級数も相性が良いです。解の構造は母関数の部分分数分解から説明できます。→ 別記事にて解説しました
