Editorial - Exawizards Programming Contest 2021（AtCoder Beginner Contest 222）

Contest Duration: 2021-10-09 21:00:00+0900 - 2021-10-09 22:40:00+0900 (local time) (100 minutes) Back to Home

Official

H - Beautiful Binary Tree Editorial by en_translator

It can be proved that a binary tree with $0$ or $1$ written on each vertex is a beautiful binary tree of degree $N$ if and only if

the root and every leaf has $1$ written on it;
the sum of integers written on all vertices is $N$;
$0$ is not written side by side,

so hereinafter we will count the number of binary trees that satisfies the conditions above.

First let us consider DP (Dynamic Programming). For $i \gt 0$, let

$a_i :=$ (the number of beautiful binary trees of degree $i$);
$b_i :=$ (the number of beautiful binary trees such that the root has $0$ written on it and the sum of values of vertices is $i$).

By dividing into cases depending on the sum of values on vertices in the subtree under a children, let us find the relationships between $a$ and $b$. First, $b_i$ can be represented with $a$ as

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

(where the former term correspond to the case where it has one children, and the latter to the case where it has two.)

Similarly, $a_i$ can be expressed with $a,b$ as

\[ \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} \]

If we compute $a_i$ and $b_i$ with these expressions in the increasing order of $i$, the problem can be solved in a total of $\mathrm{O}(N^2)$ time. (With a combination with divide-and-conquer FFT (Fast Fourier Transform), it can be sped up to $\mathrm{O}(N \log^2 N)$ time.)

Generating Function

For even better optimization, we introduce generating function. A generating function refers to a formal power series that stores information about a sequence.

A formal power series is a generalization of polynomials whose number of terms is not necessarily finite.

For example, for a sequence $a_1,a_2,\dots$, the ordinary generating function $A(x)$ of $a_n$ is defined as

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

Also, we define a notation $\lbrack x^n \rbrack$ of extracting a coefficient of a formal power series as

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

Let $A(x)$ and $B(x)$ be the generating functions of $a_i$ and $b_i$, respectively. With an aid of generating functions, the aforementioned recurrence relations can be concisely expressed as a convolution. First, $B(x)$ is expressed with $A(x)$ as

\[ \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} \]

(For simplicity, we adopt a shorthand like $B(x) \to B$.) Similarly, $A(x)$ can be written as

\[ \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} \]

In the explanation above, we found $A(x)$ through the recurrence relations, but more simply, it can be found as “the child’s generating function is $(1 + A + B)$, so $A = x (1 + A + B)^2$.”

Therefore, the desired answer can be obtained as the $N$-th coefficient of $A(x)$ satisfying

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

With this relations, we can obtain the answer in an $\mathrm{O}(N \log N)$ time (with Newton’s algorithm; search it if you are interested in.), but this is still not enough for the TL (Time Limit).

Lagrange inversion theorem

Lagrange inversion theorem is a strong weapon for tree-counting problems using generating functions.

Lagrange inversion theorem is a formula of, given a function $F(x)$, finding the coefficients of the inverse function $G(x)$ satisfying $G(F(x))=x$. Among several expression of the theorem, we will show the one that is often used for combinatorics.

Assume that two formal power series $F(x)$ and $G(x)$ satisfies $F(G(x)) = x$. If $\lbrack x^0 \rbrack F(x)=\lbrack x^0 \rbrack G(x)=0$ and $\lbrack x^1 \rbrack F(x) \neq 0, \lbrack x^1 \rbrack G(x) \neq 0$, then

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

Proof

Hereinafter, we use formal Laurent series (a formal power series with finite number of negative-power terms). First we prove the following Lemma.

Lemma (Formal residue)

For any formal power series $F(x)$ such that $F(0)=0, F'(0) \neq 0$, and any integer $k \in \Z$, the following equation holds. $$\lbrack x^{-1} \rbrack F'(x) F^k (x) = \lbrack k = -1\rbrack$$ (Proof of the lemma)

If $k \neq -1$, then it is obvious because $F'(x) F^k (x) = \frac{1}{k+1}\left(F(x)^{k+1}\right)'$.
If $k = -1$, let $F(x) = \sum_{i>0} a_i x^i$, then $$\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}$$, so the coefficient of $(-1)$-th degree for $k = -1$ is $1$.$\Box$

By the lemma, for $F(x)$ and $G(x)$ satisfying $F(G(x)) = x, F(0)=0, F'(0) \neq 0$, $$F(G(x)) = x.$$ By differentiating both hand sides, $$F'(G) \cdot G' = 1.$$ By decomposing the $F'(G)$ part by coefficients, $$\sum_{i} i (\lbrack x^i \rbrack F (x)) G^{i-1} G' = 1.$$ By multiplying $G^{-n}$ to both hand sides, $$\sum_{i} i (\lbrack x^i \rbrack F (x)) G^{i-1-n} G' = G^{-n}.$$ By extracting $\lbrack x^{-1} \rbrack$ from both hand sides, $$\lbrack x^{-1} \rbrack \sum_{i} i (\lbrack x^i \rbrack F (x)) G^{i-1-n} G' = \lbrack x^{-1} \rbrack G^{-n}$$ By the lemma, $$n \lbrack x^n \rbrack F = \lbrack x^{-1} \rbrack G^{-n}$$ The desired equation can be transformed from this equation.

Example: Catalan number

We will first find the general term of Catalan numbers, $1,1,2,5,14,\dots$, with the inversion theorem.

Let $C(x)$ be the generating function of Catalan numbers $c_i$. Let $F(x)=C(x)-1$, then the relations

\[F(x) = x(F(x) + 1)^2\]
holds (which derives from the recurrence relations), so for

\[G(x) = \frac{x}{(x+1)^2},\]
$G(F(x)) = x$ holds. Therefore

\[ \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} \]
thus, with $c_0 = 1$, we could obtain

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

Now let us adopt the inversion formula to the original problem. Again, what we want is the $N$-th coefficient of $F(x)$ that satisfies

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

Let

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

then $G(x)$ is the inverse function of $F(x)$, so by assigning it in the Lagrange inversion theorem

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

we can obtain

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

All that left is to find the $(N-1)$-th coefficient of $U(x) := T(x)^m$ for $T(x) := 1+3x+x^2$ and $m := 2N$.
At a glance, it seems that computing $U$ requires $\mathrm{O}(N \log N)$ time, but with a technique using differentials, we can find the coefficients of $U$ in an $\mathrm{O}(N)$ time.

First, by the relations between $T$ and $U$, we obtain

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

Let $u_k := \lbrack x^k \rbrack U(x)$, and by comparing the $(k-1)$-th coefficient, we obtain

left hand side: $k u_k + 3 (k-1) u_{k-1} + (k-2) u_{k-2}$
right hand side: $m(3 u_{k-1} + 2 u_{k-2})$

so the recurrence relations

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

hold. Therefore, starting from the initial term $u_0 = 1$, and using the recurrence relations and the enumeration of inverses, we can compute $u_{N-1}$ in a linear time. Therefore, the problem has been solved in a total of $\mathrm{O}(N)$ time. (Since $u_n$ is P-recursive, high-level algorithm enables us to compute them in an $\mathrm{O}(\sqrt{N} \log N)$ time, whose detail will be appended in a few hours.)

P.S. I found in a participant’s submission that $u_k$ can also be computed with a method using $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} \]

so the $k$-th term of $(u_n)$ has been computed in an $\mathrm{O}(N)$ time.

Another solution: P-recursive

There is an alternative way of using a property of the sequence called P-recursive. A sequence $a_n$ is said to be P-recursive if there exists polynomials on $n$, $p_0 (n), p_1(n), \dots$ and $p_r(n)$, such that for all $n \geq 0$,

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

In terms of generating function, let $A(x)$ be the ordinary generative function of $a_n$, then P-recursive is defined by

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

Examples of P-recursive sequence are factorials, Catalan numbers, and Montmort Numbers.

The generating function for this problem was $F(x)$ such that

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

By transforming the expression above, we can see that there exists $Q_0(x),Q_1(x),Q_2(x)$, and $Q_3(x)$ such that

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

so we can see that $F(x)$ is a generative function of a P-recursive sequence.

We will explain how to find the $n$-th term of coefficients of P-recursive sequence. First, let $\displaystyle d:= \max_{0 \leq i \leq r}\deg (p_i(n))$.

First of all, we will compute $p_0(n),\dots,p_r(n)$. To find them, we first naively compute the first $((r + 1)(d + 2) - 2)$ terms of $a_n$, and derive a simultaneous linear equations in which the coefficients sequence is regarded as a variable. Then the simultaneous equations can be solved in a total of $\mathrm{O}((rd)^3)$ to reconstruct them.
After reconstruction, all that left is to compute $a_n$ in an $\mathrm{O}(ndr)$ time with the recurrence relations, or in an $\mathrm{O}(\sqrt{nd}(r^3 + r^2 \log(nd)))$ time using an advanced algorithm consisting of doubling, Lagrange interpolation, and convolutions.

In this problem, $n = N, d = r = 3 = \mathrm{const.}$, so the problem has been solved in a total of $\mathrm{O}(N)$ or $\mathrm{O}(\sqrt{N} \log N)$ time.

The following is a sample implementation in PyPy of the solution using the inversion theorem.

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)

H - Beautiful Binary Tree Editorial by en_translator

Generating Function

Lagrange inversion theorem

Example: Catalan number

Another solution: P-recursive

Similar problems

Related articles