kmjp's blog

競技プログラミング参加記です

AtCoder ABC #222 (エクサウィザーズプログラミングコンテスト2021) : H - Beautiful Binary Tree

ARC

知らんテクが出てきた…。
https://atcoder.jp/contests/abc222/tasks/abc222_h

問題

正整数Nに対し、以下を満たす根付き木二分木をN次の美しい二分木とする。

  • 各頂点0か1が書かれている。
  • 葉には1が書かれている。
  • 頂点を1つ選び、親または親の親に自身に書かれた数値を加算して、その後頂点に書かれた数値を0にする、という処理をN-1回行うと、根頂点に書かれた数値がN、それ以外は0とできる。

Nに対し、N次の美しい二分木は何通りか。

解法

まず美しい二分木の満たす条件を整理する。

  • 最後の処理を行う頂点は、(根頂点以外)もともと1が書かれた頂点のみであり、葉から順に行う。
  • 根頂点はもともと1が書かれている。

ここから以下を考える。
f(N) := N次の美しい二分木の組み合わせ数
g(N) := 根頂点の初期値が0で、その子がいずれもa,b次(a+b=N)の美しい二分木であるような組み合わせ数

  • \displaystyle f(N) = \sum_{a=0}^{N-1} (f(a)+g(a))*(f(N-1-a)+g(N-1-a))
  • \displaystyle g(N) = \sum_{a=0}^N g(a)*g(N-a)

なので、O(N^2)で良ければこれで求まる。

これを高速化しよう。
数列f(n)・g(n)に対応する母関数をA(x)・B(x)とすると、上記式は

  • B = 2A + A^2
  • A = x(1+A+B)^2

となる。式変形してBを消すとA = x(1+3A+A^2)^2となる。
これをラグランジュの反転公式を適用すると、
AのN次の項の係数は、(1+3x+x^2)^2N/Nの(N-1)次の項になる。

頑張ってFFTすると、この値を求めることができるもののO(Nlog^2N)かかりTLEしてしまう。
微分を取ってうまく変形すると、(1+3x+x^2)^2N/Nのn次の項の係数U(n)は

  • \displaystyle U_n = \frac{3(2N+1-n)U_{n-1}+(4N+2-n)U_{n-2}}{n}

となるのでO(N)でU(N-1)を求めればよい。

int N;
const ll mo=998244353;

ll modpow(ll a, ll n = mo-2) {
	ll r=1;a%=mo;
	while(n) r=r*((n%2)?a:1)%mo,a=a*a%mo,n>>=1;
	return r;
}

ll dp[10101011][2];

void solve() {
	int i,j,k,l,r,x,y; string s;
	
	cin>>N;
	
	dp[0][1]=1;
	for(i=1;i<=N;i++) {
		/*
		// top is 1
		for(x=0;x<=i-1;x++) (dp[i][1]+=(dp[x][1]+dp[x][0])*(dp[i-1-x][1]+dp[i-1-x][0]))%=mo;
		// top is 0
		for(x=0;x<=i;x++) (dp[i][0]+=dp[x][1]*dp[i-x][1])%=mo;
		
		*/
		dp[i][1]+=3*(2*N+1-i)*dp[i-1][1];
		if(i>=2) dp[i][1]+=(4*N+2-i)*dp[i-2][1];
		dp[i][1]=dp[i][1]%mo*modpow(i)%mo;
	}
	cout<<dp[N-1][1]*modpow(N)%mo<<endl;
	
}

まとめ

知らない知識だ…。