公式

G - Balls and Boxes 解説 by Nyaan


この問題は 母関数 あるいは 形式的冪級数(FPS) と呼ばれるテクニックへ入門するための問題です。この問題および解説を通じて母関数の概念について学べたら幸いです。

母関数とは?

G.f. is the projection of combinatorial structures. — Elegia

母関数 (generating function, G.f.) とは、簡単に言うと数列の情報を保持する関数のような形をしたものです。
無限に続く数列 \(a_0, a_1, a_2, \dots\) に対して、その母関数 \(A(x)\) を次のように定義します。

\[A(x) = a_0 + a_1 x + a_2 x^2 + \dots = \sum_{i=0}^\infty a_i x^i\]

ちなみに、母関数にはいくつか種類があります。上の定義は最も一般的な母関数の定義で、区別のためにこれを 通常型母関数(ordinary generating function, OGF) と呼ぶこともあります。

また、\(A(x)\) のように項が無限に続いてもよい多項式のことを 形式的冪級数(formal power series, FPS) と呼びます。(母関数と形式的冪級数は競技プログラミングでは混同されがちですが、このような定義の違いがあります)。また、形式的冪級数の係数を取り出す記号として \(\lbrack x^n \rbrack\) を次のように定義します。

\[\lbrack x^n \rbrack A(x) := a_n\]

(読み飛ばしても問題ない補足) 後々引っ掛かりそうな話題について簡単に触れておきます。厳密に述べると、形式的冪級数は、関数、すなわち \(x\) に値を代入する対象 ではありません。簡単に言うと、\(x\) は単なる記号です。といっても無意味な記号ではなく、掛け算が定義されていれば \(x\)\(x^n\) は意味を持って定義できます。このように形式的冪級数を扱うことで、収束半径のような解析的な議論は不要となり、単なる代数的な対象、すなわち足し算や掛け算などができるものとして捉えることができます。我々が興味があるのは関数 \(A(x)\) の値ではなく \(A(x)\) の係数 \(a_n\) の値であるため、これで何の問題も起こりません。

さて、母関数は上記のように数列からシンプルに定義されるものであることを説明しました。しかし、母関数は数列から機械的に得られるだけのある種の無機的な概念ではありません。母関数は、組合せ的な対象、すなわち数え上げるオブジェクト(例えばグラフ・文字列・操作列など) を、数えたい方法に応じて “射影” した概念として考えることができます。そして、組合せ的な構造を母関数による表現で記述することでオブジェクトを数え上げる手法は適用できる範囲が非常に広いです。この応用性の高さにより、競技プログラミングにおいて母関数を利用した数え上げは多くの人が知る手法となりました。

次の章では、母関数がある種の “射影” である理由、すなわち数え上げるオブジェクトおよびその集まりと母関数がどのように密接に関わっているのかを説明します。

通常型母関数と組合せ的対象の関係

この章は ABC230 解説 を一部改変したものです。

わかりやすさのため、この章の説明では数列・母関数・集合を次に挙げる文字種を使用した表記で統一します。

  • 数列 \(\to\) 小文字斜体 (\(a, b, c, \dots\))
  • 母関数 \(\to\) 大文字斜体 (\(A, B, C, \dots\))
  • 集合 \(\to\) 大文字筆記体 (\(\mathcal{A}, \mathcal{B}, \mathcal{C}, \dots\) )

数え上げの問題では、必ず数え上げる “モノ” (グラフ, 文字列, 操作列, …) 、およびそれらからなる “集まり” (集合) が存在します。この集合が良い性質を持っているとき、集合を母関数に置き換えて考えることができます。

例としてまずは次の簡単な問題を考えてみましょう。

高橋君は弁当を買うことにしました。
食堂 A では \(i\) 円の弁当が \(a_i\) 個売られています。
食堂 B では \(j\) 円の弁当が \(b_j\) 個売られています。

(1) 高橋君は A か B の いずれか で弁当を \(1\) 個買うことにしました。どのような組合せが考えられますか?
(2) 高橋君は A と B の 両方 で弁当を \(1\) 個ずつ買うことにしました。どのような組合せが考えられますか?
(3) 「\(n\) 円かかる弁当の買い方の通り数 \(a_n\) 」の母関数を考えます。(1), (2) を母関数で表してください。

たとえば A で \(1\) 円と \(2\) 円の弁当が \(1\) 個ずつ、B で \(1\) 円と \(3\) 円の弁当が \(1\) 個ずつ場合を例として定式化してみましょう。
食堂 A の弁当を \(1\)、食堂 B の弁当を \(1^\ast\) のように表します。すると、食堂 A, B の弁当からなる集合 \(\mathcal{A}, \mathcal{B}\)

\[ \mathcal{A} = \lbrace 1, 2 \rbrace , \mathcal{B} = \lbrace 1^\ast , 3^\ast \rbrace \]

となります。
(1), (2) で数え上げる対象も同様に集合で表していきましょう。
(1) で問われている操作は交わりを持たない集合同士をマージする操作で、(2) で問われているのは集合から \(1\) 個ずつ取り出して組み合わせる操作です。
これらの操作は集合に対する広義の直和(非交和)・直積とみなすことができるので、(1) の操作を \(+\), (2) の操作を \(\times\) で表すことにしましょう。
たとえば上の例だと \(\mathcal{A} + \mathcal{B}, \mathcal{A} \times \mathcal{B}\) は次のようになります。

\[ \begin{aligned} \mathcal{A} + \mathcal{B} &= \lbrace 1,2,1^\ast, 3^\ast \rbrace \\ \mathcal{A} \times \mathcal{B} &= \lbrace 1, 2 \rbrace \times \lbrace 1^\ast , 3^\ast \rbrace \\ &= \lbrace (1,1^\ast), (1,3^\ast), (2,1^\ast), (2,3^\ast) \rbrace \end{aligned} \]

さて、次はこの集合を母関数に置き換えてみましょう。

\(\mathcal{A}, \mathcal{B}\) の母関数 \(A(x), B(x)\)

\[ A(x) = \sum_i a_i x^i, B(x) = \sum_j b_j x^j\]

となります。 (1), (2) に対応する母関数は、実は \(A(x)\)\(B(x)\) を用いて記述することができます。
まずは \(\mathcal{A} + \mathcal{B}\) です。集合 \(\mathcal{A}, \mathcal{B}\) が背反ならば、

  • (\(\mathcal{A}\) に含まれる \(n\) 円の買い方) + (\(\mathcal{B}\) に含まれる買い方) = (\(\mathcal{A} + \mathcal{B}\) に含まれる\(n\) 円の買い方)

が成り立つので、 (1) で \(i\) 円使う買い方を \(c_i\) としたとき

\[ c_i = a_i + b_i\]

が成り立ちます。上の式から (1) の母関数 \(C(x)\)

\[ \begin{aligned} C(x) &= \sum_{i} (a_i + b_i) x^i \\ &= A(x) + B(x) \end{aligned} \]

であることがわかります。

次に \(\mathcal{A} \times \mathcal{B}\) です。(2) で \(i\) 円使う買い方を \(d_i\) としたとき、

\[ d_k = \sum_{i + j = k} a_i b_j\]

となるので、

\[ \begin{aligned} D(x) &= \sum_k \left(\sum_{i+j = k} a_i b_j\right) x^k \\ &= \sum_i a_i x^i \sum_j b_j x^j \\ &= A(x) \times B(x) \end{aligned} \]

となることが分かります。結局まとめると

\[\mathcal{C} = \mathcal{A} + \mathcal{B} \iff C(x) = A(x) + B(x)\]

\[\mathcal{D} = \mathcal{A} \times \mathcal{B} \iff D(x) = A(x) \times B(x)\]

となることがわかります。たとえば \(\mathcal{A} = \lbrace 1, 2 \rbrace ,\mathcal{B} = \lbrace 1^\ast , 3^\ast \rbrace\) の場合、

\[ A(x) = x + x^2, B(x) = x + x^3\]

\[C(x) = A(x) + B(x) = 2x + x^2 + x^3\]

\[D(x) = A(x) \times B(x) = x^2 + x^3 + x^4 + x^5\]

となり、さきほどの例示と比較すると \(C(x), D(x)\)\(\mathcal{A} + \mathcal{B}, \mathcal{A} \times \mathcal{B}\) の母関数になっていることが確認できます。

上の例ではさすがに当たり前すぎて何の役に立つのかよくわからないと思うので、もう少し使い道のありそうな例で考えていきましょう。

高橋君は階段を \(1\) ステップに \(1\) 段または \(2\) 段ずつ進んで登ることにしました。\(N = 0,1,\dots, n\) について、\(N\) 段の階段の登り方は何通りありますか?

「階段を登る方法」を数列として表すと、条件を満たす数列は \(1, 2\) からなる長さ \(0\) 以上の数列になります。よって数え上げる対象からなる集合 \(\mathcal{F}\)

\[\mathcal{F} = \lbrace \emptyset, (1), (2), (1, 1), (1, 2),(2,1), (2,2), \dots \rbrace\]

のようになります。
ここで、\(\mathcal{F}\) の要素をはじめに \(1\) 段登ったか \(2\) 段登ったかの \(2\) 通りに分類しましょう。すると \(\mathcal{F}\) はある種の再帰的な性質を持っているので

\[ \begin{aligned} &\mathcal{F} \\ &= \lbrace \emptyset \rbrace + \lbrace (1) \rbrace \times \lbrace \emptyset, (1), (2),\dots, \rbrace + \lbrace (2) \rbrace \times \lbrace \emptyset, (1), (2),\dots, \rbrace \\ &= \lbrace \emptyset \rbrace + \lbrace (1) \rbrace \times \mathcal{F} + \lbrace (2) \rbrace \times \mathcal{F} \end{aligned} \]

と記述することができます。

さて、\(\mathcal{F}\) を「段数」に注目した母関数に置き換えてみましょう。段数を \(x\) の指数部分に置き換えると、\(\emptyset\)\(x^0 = 1\)\((1)\)\(x\)\((2)\)\(x^2\) に置き換えられるので、\(\mathcal{F}\) に対応する母関数 \(F(x)\)

\[ \begin{aligned} &F(x) = 1 + xF(x) + x^2 F(x) \\ &\to F(x) = \frac{1}{1 - x - x^2} \end{aligned} \]

であるとわかります。実際、 \(\frac{1}{1 - x - x^2}\) を級数展開すると

\[F(x) = 1 + x + 2x^2 + 3x^3 + 5x^4 + 8x^5 + \dots\]

となり、求める数列 \((f_n)\) はフィボナッチ数列

\[ (f_n) = 1, 1, 2, 3, 5, 8, 13, 21, \dots \]

と一致することが確認できます。

2 つの例が意味することとして、結局、母関数とは、数え上げる対象からなる「集合」から、必要な情報 (ここでは価格や段数) 以外を捨象したものになるということが言えます。
そして、数え上げる対象が良い性質を持っている場合、対象からなる集合の大きさは母関数に対する加算や乗算などの初等的な演算で記述できて、計算できる形に持ち込むことができます。このような操作を機械的に行うことができるのが母関数を用いる 1 つの利点となります。

発展:組合せ的な “構造” と関数の合成

この章は ABC230 解説 を一部改変したものです。

前章で出てきたのは \(+\)\(\times\) という基本的な演算の \(2\) つのみでしたが、次はもう少し大きな “パーツ” を扱ってみましょう。

数え上げる対象からなる集合は、ある種の組合せ的な “構造” の上に載っていることが多いです。(例:列, 円環, 集合, …) このような構造は、集合から集合への写像 \(\mathrm{C}(\mathcal{A})\) として扱うことができます。

ここで特筆すべき点としては、\(\mathrm{C}\) がどのような関数か分かっていれば、 \(\mathcal{A}\)\(\mathrm{C}\) に代入することでただちに \(\mathrm{C}(\mathcal{A})\) の集合を描写することができる点にあります。
これを応用すると、数え上げたい対象が複雑な構造をしていても、それがどのような集合の合成で出来ているかを把握することができれば、いくつかの簡単な部品に “分解” して考察を大きく単純化・機械化することができます。

“構造” の簡単な例として、集合 \(\mathcal{A}\) に対して「 \(\mathcal{A}\) の要素からなる長さ \(0\) 以上の列」に対応する構造 \(\mathrm{SEQ}(\mathcal{A})\) を考えてみましょう。列の長さが \(i\) のときは \(\mathcal{A}^i\) になることから、 \(\mathrm{SEQ}(\mathcal{A})\)\(\mathcal{E}\) を空の列を意味する集合として次のように表すことができます。

\[\mathrm{SEQ}(\mathcal{A}) = \mathcal{E} + \mathcal{A} + \mathcal{A} \times \mathcal{A} + \mathcal{A} \times \mathcal{A} \times \mathcal{A} + \dots = \sum_{i=0}^\infty \mathcal{A}^i\]

これを母関数の世界で言い換えると、無限級数の公式から

\[S(x) = 1 + A(x) + A(x)^2 + \dots = \sum_{i=0}^\infty A(x) = \frac{1}{1 - A(x)}\]

と変形することができます。

構造 \(\mathrm{SEQ}(\mathcal{A})\) をさっそく問題に適用していきましょう。さきほどのフィボナッチ数列の例も \(\mathrm{SEQ}(\mathcal{A})\) に代入することで機械的に求めることができます。「階段の登り方」は長さ \(0\) 以上の登り方の列とみなせて、「階段を \(1\) 段か \(2\) 段登る」という操作 \(\mathcal{A}\) の母関数は

\[A(x) = x + x^2\]

なので、これを上式に代入することで、フィボナッチ数列の母関数 \(F(x)\) は、

\[F(x) = \frac{1}{1 - A(x)} = \frac{1}{1 - x - x^2}\]

であると直ちにわかります。

このように数え上げる対象がいくつかの “構造” の組み合わせでできているとき、それらの “構造” 一つ一つの関数さえわかっていれば、それらを組み合わせて関数を合成することで数え上げる対象の母関数をただちに求めることができます。

指数型母関数

数列の母関数を考える時、しばしば \(a_0, a_1, \dots, a_n, \dots, \) の母関数を考えると複雑な表現になってしまう一方で、\(\dfrac{a_0}{0!}, \dfrac{a_1}{1!}, \dots, \dfrac{a_n}{n!},\dots\) の母関数を考えると上手くいくことがあります。このようなことは非常に多いので、この母関数には 指数型母関数(exponential generating function, EGF) という名前がついています。すなわち次式の \(\hat{A}(x)\) が指数型母関数です。

\[\hat{A}(x) = \sum_{n=0}^\infty \frac{a_n}{n!} x^n\]

指数型母関数の係数を取り出す記号として \(\left\lbrack \dfrac{x^n}{n!} \right\rbrack\) を次のように定義します。

\[\left\lbrack \dfrac{x^n}{n!} \right\rbrack \hat{A}(x) := a_n\]

指数型母関数の重要な特徴として、指数型母関数同士の畳み込みが二項係数の重み付きの畳み込みになるという点が挙げられます。すなわち、次式が成り立ちます。

\[ \begin{aligned} &\left\lbrack \dfrac{x^k}{k!} \right\rbrack \hat{A}(x) \hat{B}(x) \\ &= \left\lbrack \dfrac{x^k}{k!} \right\rbrack \left(\sum_{i=0}^\infty \frac{a_i}{i!}x^i\right) \left(\sum_{j=0}^\infty \frac{b_j}{j!}x^j\right)\\ &= k! \sum_{i+j=k} \frac{a_i}{i!} \frac{b_j}{j!}\\ &= \sum_{i+j=k} \binom{k}{i} a_i b_j \end{aligned} \]

指数型母関数の性質を一通り説明しました。次は、指数型母関数と組合せ的な構造の関係性を見ていき、なぜ指数型母関数が良い性質を持つのかを説明します。

指数型母関数と組合せ的対象の関係

先に述べた通り、通常型母関数の加法と乗法は集合の非交和と直積に対応するものでした。 指数型母関数の加法と乗法を同様の方法で解釈すると、実は ラベル付きオブジェクトの集合 の非交和と直積に対応するものになっています。どういうことなのか詳しく見てみましょう。

まず、ラベル付きオブジェクト \(a\) を次のように定義します。

  • \(a\)\(\vert a \vert\) 個の “容器(例えばグラフの頂点など)” を持つオブジェクトである。\(a\) の容器には \(1, 2, \dots, |a|\) というラベルが張られている。

そして、ラベル付きオブジェクト \(a, b\) を結合して新たなオブジェクトの集合を作る操作 \(a \otimes b\) を次のように定義します。

  • ラベル \(1, 2, \dots, |a| + |b|\) が貼られた \(|a|+|b|\) 個の容器を大きさ \(|a|\) の集合と大きさ \(|b|\) の集合に分割する方法は \(\binom{|a| + |b|}{|a|}\) 通りある。それぞれの分割に対して、前者の集合にオブジェクト \(a\) を、後者の集合にオブジェクト \(b\) をコピーしてできるオブジェクトを新たに生成する。こうしてできた \(\binom{|a| + |b|}{|a|}\) 個のオブジェクトからなる集合を \(a \otimes b\) とする。

例えば、\(a\) が赤く塗られた \(3\) 個のマス、\(b\) が青く塗られた \(1\) 個のマス、とします。すると、\(\text{R}\) を赤、\(\text{B}\) を青、文字列の \(i\) 文字目をラベル \(i\) が貼られたマスとした時、\(a \otimes b\)

\[a \otimes b = \lbrace \text{RRRB, RRBR, RBRR, BRRR}\rbrace\]

というようになります。これはラベル付きオブジェクト同士の結合として自然な定義だと思います。

さて、オブジェクト同士の結合を、オブジェクトの容器の個数に注目した母関数として表してみましょう。通常のラベル無しオブジェクトの場合は、大きさ \(1\) の集合同士の直積もまた大きさ \(1\) の集合になるため、

\[x^{|a| + |b|} = x^{|a|} \times x^{|b|}\]

という単純な形で表現することができるのでした。しかし今回は、\(1\) 個のオブジェクト同士から \(\binom{|a| + |b|}{|a|}\) 個のオブジェクトたちが生成されるので上式ではうまくいきません。

そこで、母関数に容器の個数に応じた重みづけをして上手くいかないか考えてみることにします。つまり、サイズ \(|a|\) のオブジェクトに \(w(|a|)\) の重みをつけて \(w(|a|) x^{|a|}\) という母関数を対応させることを考えます。

\(w\) としてどのような関数がふさわしいかを考えてみましょう。条件を満たす関数 \(w\)

\[\binom{|a| + |b|}{|a|} w(|a| + |b|) x^{|a|+|b|} = w(|a|) x^{|a|} \times w(|b|) x^{|b|}\]

を満たします。ここで \(\binom{|a| + |b|}{|a|} = \frac{(|a|+|b|)!}{|a|!|b|!}\) であることを踏まえると、 \(w(n) = \frac{1}{n!}\) が任意の状況において条件を満たすことが証明できます。
つまり、ラベル付きオブジェクト \(a\) に対して

\[\frac{x^{|a|}}{|a|!}\]

を対応させることで、

\[\binom{|a| + |b|}{|a|} = |a \otimes b|\]

というオブジェクトの合成操作を母関数で表すと

\[\binom{|a| + |b|}{|a|} \frac{x^{|a|+|b|}}{(|a|+|b|)!} = \frac{x^{|a|}}{|a|!} \times \frac{x^{|b|}}{|b|!}\]

という式になり、上式は恒等式なので全てが上手くいくという話です。

ここまでは \(1\) 個のラベル付きオブジェクト同士の合成について考えていましたが、これはラベル付きオブジェクトの集合同士の直積に拡張することが出来ます。

ラベル付きオブジェクトの集合 \(\mathcal{A}\) の母関数を次のように定義します。

\[A(x) = \sum_{a \in \mathcal{A}} \frac{x^{|a|}}{|a|!} = \sum_{n=0}^\infty \frac{a_n}{n!} x^n\]

そして、2 つの排反なラベル付きオブジェクトの集合 \(\mathcal{A}, \mathcal{B}\) に対して

\[\mathcal{A} + \mathcal{B} = \mathcal{A} \cup \mathcal{B}\]

\[\mathcal{A} \times \mathcal{B} = \bigcup_{a \in \mathcal{A}, b \in \mathcal{B}} a \otimes b\]

と定義すると、

\[\mathcal{C} = \mathcal{A} + \mathcal{B} \iff C(x) = A(x) + B(x)\]

\[\mathcal{D} = \mathcal{A} \times \mathcal{B} \iff D(x) = A(x) \times B(x)\]

が成り立つことが確認できます。
このように、通常型母関数と指数型母関数は一見すると別物ですが、実は加算と乗算が集合の非交和と直積に対応しているという点で同種の概念であるということが確認できました。

以上の内容が母関数の組合せ的な背景を説明したものになります。この解説を通じて母関数への理解が深まったら幸いです。

今回の問題の答え

ここまでの説明を踏まえたうえで今回の問題を見てみます。すると、問題 1 は(ラベル無し)オブジェクトの数え上げ、問題 2 はラベル付きオブジェクトの数え上げです。よって問題 1 は通常型母関数を、問題 2 は指数型母関数を適用すれば解けることがわかります。

具体的には、問題 1 は

\[F_1(x) = \sum_{i=0}^\infty x^{iA}, F_2(x) = \sum_{i=0}^\infty x^{iB}, F_3(x) = \sum_{i=0}^\infty x^{iC}\]

とした時の

\[\lbrack x^N \rbrack F_1(x) F_2(x) F_3(x)\]

が答えで、問題 2 は

\[G_1(x) = \sum_{i=0}^\infty \frac{x^{iA}}{(iA)!}, G_2(x) = \sum_{i=0}^\infty \frac{x^{iB}}{(iB)!}, G_3(x) = \sum_{i=0}^\infty \frac{x^{iC}}{(iC)!}\]

とした時の

\[\left\lbrack \frac{x^N}{N!} \right\rbrack G_1(x) G_2(x) G_3(x)\]

が答えです。

母関数同士の \(\text{mod }998244353\) 上の畳み込みは atcoder::convolution などの畳み込みライブラリを用いれば母関数の項数を \(n\) として \(\mathrm{O}(n \log n)\) で行うことが出来ます。よって、\(F_1(x), F_2(x), F_3(x), G_1(x), G_2(x), G_3(x)\) の先頭 \(N+1\) 項を計算して畳み込めば今回の問題を \(\mathrm{O}(N \log N)\) で解くことが出来て、これは十分高速です。

  • 実装例(C++)
#include <iostream>
#include <vector>
using namespace std;

#include "atcoder/convolution.hpp"
#include "atcoder/modint.hpp"
using mint = atcoder::modint998244353;

int main() {
  int N, A, B, C;
  cin >> N >> A >> B >> C;
  vector<mint> fac(N + 1), invfac(N + 1);
  fac[0] = 1;
  for (int i = 1; i <= N; i++) fac[i] = fac[i - 1] * i;
  invfac[N] = fac[N].inv();
  for (int i = N; i >= 1; i--) invfac[i - 1] = invfac[i] * i;

  {
    vector<mint> F1(N + 1), F2(N + 1), F3(N + 1);
    for (int i = 0; i <= N; i++) {
      if (i % A == 0) F1[i] = 1;
      if (i % B == 0) F2[i] = 1;
      if (i % C == 0) F3[i] = 1;
    }
    auto mid = atcoder::convolution(F1, F2);
    auto F = atcoder::convolution(mid, F3);
    cout << F[N].val() << "\n";
  }

  {
    vector<mint> G1(N + 1), G2(N + 1), G3(N + 1);
    for (int i = 0; i <= N; i++) {
      if (i % A == 0) G1[i] = invfac[i];
      if (i % B == 0) G2[i] = invfac[i];
      if (i % C == 0) G3[i] = invfac[i];
    }
    auto mid = atcoder::convolution(G1, G2);
    auto G = atcoder::convolution(mid, G3);
    cout << (G[N] * fac[N]).val() << "\n";
  }
}

投稿日時:
最終更新: