\(0,1\) が書きこまれた二分木のうち \(N\) 次の美しい二分木である必要十分条件は

• 根・葉に書きこまれた数が \(1\)
• 頂点に書きこまれた数の和が \(N\)
• \(0\) が隣り合わない

であることが証明できるので、以下では上の条件を満たす二分木の個数を数えていきます。

さて、まずは DP を考えてみましょう。\(i \gt 0\) である \(i\) に対して

• \(a_i :=\) ( \(i\) 次の美しい二分木の個数 )
• \(b_i :=\) ( 美しい二分木の部分木のうち、根が \(0\) で頂点の値の和が \(i\) である木の個数 )

とおきます。 子についた部分木の頂点の値の和で場合分けを行うことで、 \(a\) と \(b\) の関係式を求めてみましょう。まず \(b_i\) は \(a\) を用いて

\[ b_i = 2 a_i + \sum_{0 \lt j \lt i} a_j a_{i-j}\]

と表せます。( 前半は子が \(1\) 個の場合、後者が子が \(2\) 個の場合です。)

同様に \(a_i\) は \(a,b\) を用いて

\[ \begin{aligned} a_1 &= 1 \\ a_i &= 2 (a_{i-1} + b_{i-1}) + \sum_{0 \lt j \lt i-1} (a_j + b_j) (a_{i-1-j} + b_{i-1-j}) & (i \gt 1) \end{aligned} \]

と表せます。
この式を用いて \(a_i,b_i\) を昇順に計算すれば \(\mathrm{O}(N^2)\) でこの問題を解くことができます。 (分割統治 FFT と組み合わせれば \(\mathrm{O}(N \log^2 N)\) に高速化できます。)

母関数

さらなる高速化のために 母関数 を導入します。母関数とは、数列に関する情報を保持する 形式的冪級数 のことをいいます。

• 形式的冪級数とは多項式を一般化したもので、項が有限個でなくてもよい多項式のことを言います。

たとえば、数列 \(a_1,a_2,\dots\) に対して、 \(a_n\) の 通常型母関数 \(A(x)\) は

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

として定義されます。また、形式的冪級数の係数を取り出す記号として \(\lbrack x^n \rbrack\) を

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

のように定義します。

\(a_i, b_i\) の母関数をそれぞれ \(A(x), B(x)\) としましょう。母関数を持ち出すと、前述の漸化式は畳み込みを用いてきれいに表現することができます。まず、 \(B(x)\) は \(A(x)\) を使うと

\[ \begin{aligned} b_i &= 2 a_i + \sum_{0 \lt j \lt i} a_j a_{i-j} \\ &\iff B = 2 A + A^2 \end{aligned} \]

となります。(簡単のために \(B(x) \to B\) のように略記しています。) 同様に \(A(x)\) は

\[ \begin{aligned} a_1 &= 1 \\ a_i &= 2 (a_{i-1} + b_{i-1}) + \sum_{0 \lt j \lt i-1} (a_j + b_j) (a_{i-1-j} + b_{i-1-j}) & (i \gt 1) \\ &\iff A = x + 2 x A + 2 x B + x(A + B)^2 \\ &\iff A = x(1 + A + B)^2 = x(1 + 3A + A^2)^2 \end{aligned} \]

と表すことができます。

• 上の説明では漸化式を経由して \(A(x)\) を求めましたが、 「 子の母関数は \((1 + A + B)\) だから \(A = x (1 + A + B)^2\) 」という風に求めた方が簡便です。

以上より、求める答えは

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

を満たす \(A(x)\) の \(N\) 次の係数として表すことができます。この式を使うと \(\mathrm{O}(N \log N)\) で答えを得られますが (興味がある人はニュートン法で調べてみてください) これでも TL には間に合いません。

ラグランジュの反転公式

母関数を用いた木の数え上げにおける強力な武器として ラグランジュの反転公式 と呼ばれる公式があります。

ラグランジュの反転公式は関数 \(F(x)\) が与えられたときに \(G(F(x))=x\) である逆関数の係数を求める公式です。表現形はいくつかあるようですが、数え上げで多く用いる式を以下に挙げます。

形式的冪級数 \(F(x),G(x)\) は逆関数の関係にあるとする。(すなわち \(F(G(x)) = G(F(x)) = x\) が成り立つ。) 今、 \(\lbrack x^0 \rbrack F(x)=\lbrack x^0 \rbrack G(x)=0\) かつ \(\lbrack x^1 \rbrack F(x) \neq 0, \lbrack x^1 \rbrack G(x) \neq 0\) が成り立つとき、

\[\lbrack x^n \rbrack G(x) = \frac{1}{n} \lbrack x^{n-1} \rbrack \left(\frac{x}{F(x)}\right)^n\]

が成り立つ。

証明

以下では形式的ローラン級数 (有限項の負冪項を認めた形式的冪級数) が登場します。 まず、次の補題を証明します。

補題 (形式留数)

$F(0)=0, F'(0) \neq 0$ である形式的冪級数 $F(x)$ および全ての $k \in \Z$ に対して次の式が成り立つ。 $$\lbrack x^{-1} \rbrack F'(x) F^k (x) = \lbrack k = -1\rbrack$$ (補題の証明)

$k \neq -1 $ の時は $F'(x) F^k (x) = \frac{1}{k+1}\left(F(x)^{k+1}\right)'$より明らかである。
$k = -1$ の時は$F(x) = \sum_{i>0} a_i x^i$ とすると $$\frac{F'(x)}{F(x)} = x^{-1} \frac{1 + 2\frac{a_2}{a_1}x + \cdots}{1 + \frac{a_2}{a_1}x + \cdots}$$ よって $k = -1$ のときの $-1$ 次の係数は $1$ である。$\Box$

補題を利用すると、 $F(G(x)) = x, F(0)=0, F'(0) \neq 0$ を満たす $F(x),G(x)$ に対して $$F(G(x)) = x$$ 両辺を微分して $$F'(G) \cdot G' = 1$$ $F'(G)$の部分を係数ごとに分解して $$\sum_{i} i (\lbrack x^i \rbrack F (x)) G^{i-1} G' = 1$$ 両辺に $G^{-n}$ を掛けて $$\sum_{i} i (\lbrack x^i \rbrack F (x)) G^{i-1-n} G' = G^{-n}$$ 辺々 $\lbrack x^{-1} \rbrack$ を取って $$\lbrack x^{-1} \rbrack \sum_{i} i (\lbrack x^i \rbrack F (x)) G^{i-1-n} G' = \lbrack x^{-1} \rbrack G^{-n}$$ 補題を利用して $$n \lbrack x^n \rbrack F = \lbrack x^{-1} \rbrack G^{-n}$$ この式を変形すれば題意の式を得ることができる。

例題：カタラン数

カタラン数 \(1,1,2,5,14,\dots\) の一般項を反転公式を使って求めてみましょう。

カタラン数 \(c_i\) の母関数を \(C(x)\)とします。 \(F(x)=C(x)-1\)とおくと

\[F(x) = x(F(x) + 1)^2\]

という関係式が成り立つので ( 漸化式を立てると従います ) 、

\[G(x) = \frac{x}{(x+1)^2}\]

とおくと\(G(F(x)) = x\) が成り立ちます。よって

\[ \begin{aligned} [x^n]F(x) &= \frac{1}{n} \lbrack x^{n-1} \rbrack \left(\frac{x}{G(x)}\right)^n \\ &= \frac{1}{n} \lbrack x^{n-1} \rbrack (x + 1)^{2n} \\ &= \frac{1}{n} \binom{2n}{n-1} = \frac{1}{n+1} \binom{2n}{n} \end{aligned} \]

より、 \(c_0 = 1\) と合わせて

\[c_n = \frac{1}{n+1} \binom{2n}{n}\]

を得ることができました。

さて、反転公式を元の問題に適用していきましょう。今、求めたいものは

\[F(x) = x(1 + 3F(x) + F(x)^2)^2\]

を満たす \(F(x)\) の \(N\) 次の係数でした。

\[G(x) := \frac{x}{(1+3x+x^2)^2}\]

とおくと \(G(x)\) は \(F(x)\) の逆関数になっているので、ラグランジュの反転公式

\[ \lbrack x^n \rbrack F(x) = \frac{1}{n} \lbrack x^{n-1} \rbrack \left( \frac{x}{G(x)} \right)^n \]

に代入して

\[ \lbrack x^N \rbrack F(x) = \frac{1}{N} [x^{N-1}] (1+3x+x^2)^{2N} \]

を得ることができました。

あとは多項式 \(T(x) := 1+3x+x^2\) と正の整数 \(m := 2N\) に対して \(U(x) := T(x)^m\) の \(N-1\) 次の係数が計算できれば良いです。
\(U\) の計算は一見すると \(\mathrm{O}(N \log N)\) かかりそうですが、微分を利用したテクニックを使用すれば \(U\) の係数を \(\mathrm{O}(N)\) で求めることができます。

まず、 \(T\) と \(U\) の関係式から

\[U' = m T^{m-1} T' \to T U' = m T' U\]

を得ます。\(u_k := \lbrack x^k \rbrack U(x)\) のようにおいて \(k - 1\) 次の係数を比較すると

• 左辺：\(k u_k + 3 (k-1) u_{k-1} + (k-2) u_{k-2}\)
• 右辺：\(m(3 u_{k-1} + 2 u_{k-2})\)

になるので、漸化式

\[u_k = \frac{3(m + 1 - k)u_{k-1} + (2m + 2 - k)u_{k-2}}{k}\]

が成り立ちます。よって初項 \(u_0 = 1\) からスタートして漸化式 + 逆元の列挙を利用すれば \(u_{N-1}\) を線形時間で計算することができます。以上よりこの問題を \(\mathrm{O}(N)\) で解くことができました。(詳細は後述しますが、 \(u_n\) は P-recursive なので高度なアルゴリズムを利用すれば \(\mathrm{O}(\sqrt{N} \log N)\) で計算することもできます。)

追記: 参加者の方の提出を読んで気づきましたが、二項係数を利用した方法でも \(u_k\) を計算できます。

\[ \begin{aligned} u_k &= \lbrack x^k \rbrack (1+3x+x^2)^{2N} \\ &= \lbrack x^k \rbrack ( (1+3x)+x^2)^{2N} \\ &=\lbrack x^k \rbrack \sum_{0 \leq j \leq 2N} \binom{2N}{j}x^{2j} (1+3x)^{2N-j} \\ &=\sum_{0 \leq j \leq 2N} \binom{2N}{j}\lbrack x^{k-2j} \rbrack (1+3x)^{2N-j} \\ &=\sum_{0 \leq j \leq 2N} \binom{2N}{j} \binom{2N-j}{k-2j} 3^{k-2j} \end{aligned} \]

より \((u_n)\) の第 \(k\) 項を \(\mathrm{O}(N)\) で求められました。

別解：P-recursive

別解として P-recursive と呼ばれる数列の性質を使う方法もあります。数列 \(a_n\) が P-recursive であるとは、\(n\) の多項式 \(p_0 (n), p_1(n), \dots, p_r(n)\) が存在して、すべての \(n \geq 0\) に対して

\[ p_0(n) a_n + p_1(n) a_{n+1} + \dots + p_r(n) a_{n+r} = 0\]

が成り立つことを言います。これは母関数に言い換えると、\(a_n\) の通常型母関数を \(A(x)\) としたとき、ある多項式 \(Q_0 (x), Q_1(x), \dots, Q_n(x)\) が存在して

\[ Q_0(x) A(x) + Q_1(x) A'(x) + \dots + Q_r(x) A^{(r)}(x) = 0\]

が成り立つことをいいます。

P-recursive な数列の例としては、階乗, カタラン数, Montmort Number(攪乱順列) などが挙げられます。

この問題で求める母関数は

\[F = x(1 + 3F + F^2)^2\]

を満たす \(F(x)\) でした。上式は変形すると

\[ Q_0(x)F + Q_1(x) F' + Q_2(x) F'' + Q_3(x) F''' = 0\]

を満たす \(Q_0(x),Q_1(x),Q_2(x),Q_3(x)\) が存在するとわかるので、 \(F(x)\) は P-recursive な数列の母関数であることがわかります。

P-recursive な数列の係数の第 \(n\) 項を求める方法を説明します。まず \(d\) を \(\displaystyle d:= \max_{0 \leq i \leq r}\deg (p_i(n))\)と置きます。

まずはじめに \(p_0(n),\dots,p_r(n)\) を求めます。 これは \(a_n\) の前から \((r + 1)(d + 2) - 2\) 項を愚直に計算して、計算結果から係数列を変数とみなした 1 次連立方程式を導出します。あとは連立方程式を掃き出し法で解けば \(\mathrm{O}((rd)^3)\) で復元することができます。
そして、復元したあとは \(a_n\) を漸化式を利用して \(\mathrm{O}(ndr)\) 、あるいはダブリングとラグランジュ補間と畳み込みを組み合わせた高度なアルゴリズムを利用して \(\mathrm{O}(\sqrt{nd}(r^3 + r^2 \log(nd)))\) で計算できます。

この問題では \(n = N, d = r = 3 = \mathrm{const.}\) なので、この問題を \(\mathrm{O}(N)\) または \(\mathrm{O}(\sqrt{N} \log N)\) で解くことができました。

PyPy による反転公式を利用した解法の実装例は次の通りです。

N = int(input())
mod, m, u, inv = 998244353, N * 2, [1] + [0] * (N + 10), [1] * (N + 10)
for k in range(1, N):
  u[k] = (u[k - 1] * 3 * (m + 1 - k) + u[k - 2] * (2 * m + 2 - k)) % mod * inv[k] % mod
  inv[k + 1] = mod - inv[mod % (k + 1)] * (mod // (k + 1)) % mod
print(u[N - 1] * inv[N] % mod)

類題

この問題と関連性が深い問題はコンテストだと CF Div.1 を中心に何問か出題されています。興味がある方はこちらも解いてみるとよいと思います。

• CF 438E
  • プログラミングコンテストで母関数を利用して二分木を数え上げる問題といえばこの問題が有名です。 ニュートン法 の練習やライブラリの verify 問題としても適切な難易度です。
• CF 1349F2
  • 母関数を利用した多くのテクニックを発明した Elegia 氏により出題された問題で、考察部分・計算部分ともに非常に難しいです。考察の後半にラグランジュの反転公式を利用したテクニックが登場します。
• TCO19 *** (ネタバレ回避)
  • 想定解は正統的な数え上げですが、参加者の一部は P-recursive で解いたらしい？問題です。非想定解法として P-recursive で解ける問題はこの問題に限らず少なくない印象があります。
• CF 1479E
  • 想定解では P-recursive な数列を \(\mathrm{O}(\sqrt{N} \log N)\) で計算するアルゴリズムを要求している問題です。前半の マルチンゲール を利用した考察部分も含めて、最近の数え上げの流行を反映した問題と言えると思います。

関連記事

形式的冪級数を学ぶ上でおすすめの文章やこの問題に関連する記事をいくつか紹介します。

• maspy, 形式的べき級数解説
  • 日本人が形式的冪級数を学びたい場合、真っ先に読むべきホームページです。形式的冪級数を平易な書き方で説明していて非常に読みやすく、練習問題も多く載っていて初心者の方が勉強しやすいつくりになっています。
• 37zigen, 指数型母関数入門
  • この解説では触れませんでしたが、母関数には通常型母関数のほかに 指数型母関数 と呼ばれる重要な概念があります。この記事は指数型母関数を解説した記事で、出題例が多く載っていて参考になると思います。
• zscoder, Generating Functions in Competitive Programming (Part 2)
  • 競技プログラミングでの母関数を利用した問題を丁寧に解説している記事です。 Part 1 も勉強になりますが、 Part 2 ではラグランジュの反転公式を利用した問題の解説が多く載っていてこちらの方がこの問題と関連性が深いです。
• Retired_MiFaFaOvO, A problem collection of ODE and differential technique
  • AtCoder 金冠である apiad 氏による母関数のテクニックに関する記事です。この問題に出てくる微分を利用したテクニックや P-recursive に関する話が載っています。
• Elegia, Athekatelan, 信息学竞赛中的生成函数计算理论框架
  • 題名を翻訳すると “プログラミングコンテストにおける母関数の計算理論の構造” のようになると思います。ここ 1,2 年 で中国から発信された最先端のテクニックがまとめられていて、ラグランジュの反転公式や微分を利用したテクニックもいくつか載っています。
• noshi91, polynomial_matrix_prod.rs
  • P-recursive な数列の第 \(N\) 項を \(\mathrm{O}(\sqrt{N} \log N)\) で計算するために必要なアルゴリズムの実装例です。多項式を係数とする行列 \(M(n)\) に対して \( M(N) \times M(N-1) \times \dots \times M(2) \times M(1)\) のような式を高速に計算することができれば P-recursive を高速に計算することができて、 noshi さんのライブラリはこれを実装したものになっています。