Xmas Contest 2020 H: Hierarchical Phylogeny 解説

https://atcoder.jp/contests/xmascon20/tasks/xmascon20_h

問題概要

頂点に $\{0, \ldots, N-1\}$ の空でない部分集合のラベルがついた根付き木であって，各頂点が

子が $0$ 個で，ラベルが指定された部分集合 ( $L$ ) のいずれか
子が $2$ 個で，ラベルが子のラベルの非交和

のいずれかであるようなもののうち，根のラベルが指定された部分集合 ( $R$ ) のいずれかであるようなものの個数を $\operatorname{mod} 998244353$ で求めよ．

立式

$\{0, \ldots, N-1\}$ の部分集合全体を $\mathcal{P}$ とします．入力の $L$ を $\mathcal{P}$ 上の関数として表しておきます ( $A \in \mathcal{P}$ が葉にふさわしいとき $L(A) = 1$ ，そうでないとき $L(A) = 0$ )．

入力の $R$ でいろいろ指定されていますが，結局，根のラベルを指定したときの木の個数を求められればよいです． $A \in \mathcal{P}$ に対し， $f(A)$ を条件を満たす木であって根のラベルが $A$ であるものの個数とします ( $f(\emptyset) = 0$ です)．根の子の個数が $0$ か $2$ かで場合分けして，漸化式 $f(A) = L(A) + \frac12 \sum_{B\sqcup C=A} f(B) f(C)$ が従います ( $\sqcup$ で非交和を表すことにします)．係数 $\frac12$ は子の順番を区別しないためのものです．これをそのまま DP すると $O(3^N)$ 時間になります．

subset convolution

$g, h\colon \mathcal{P} \to \mathbb{C}$ に対して，それらの subset convolution $g * h\colon \mathcal{P} \to \mathbb{C}$ は $(g * h)(A) = \sum_{B\sqcup C=A} g(B) h(C)$ で定義されます．この解説では単に積のように $g h$ と書いてしまうことにします．代数の言葉で書くと，和を各点での和，積を subset convolution とすることで $\mathcal{P}$ から $\mathbb{C}$ への関数全体が環となります．乗法の単位元は $1(\emptyset) = 1$ ， $1(A) = 0$ ( $A \ne \emptyset$ ) なる $1$ です．

subset convolution は $O(2^N N^2)$ 時間で計算できることが知られています [1]．論文を読んでいただくのがいいと思うのですが，ここでもアルゴリズムを簡単に説明しておきます．メインアイデアは $B \sqcup C = A$ が「 $B \cup C = A$ かつ $\lvert B \rvert + \lvert C \rvert = \lvert A \rvert$ 」と同値ということの利用で，集合のサイズの情報を持たせておいてから $\cup$ についての畳み込み (累積和→積→差分) をします．すなわち $g$ に対し $\hat{g}\colon \mathcal{P} \to \mathbb{C}[z]$ を $\hat{g}(A) = \sum_{B\subseteq A} g(B) z^{\lvert B\rvert}$ で定め， $\hat{h}$ も同様にし，これらの各点積 (多項式の積) を計算し $\hat{k}(A) = \hat{g}(A) \hat{h}(A)$ ，Möbius 変換で戻し $k(A) = \sum_{B\subseteq A} (-1)^{\lvert A\setminus B\rvert} \hat{k}(B)$ ，集合のサイズと次数が合っているところをとればよいです $(gh)(A) = [z^{\lvert A\rvert}] k(A)$ ．多項式の積は $\operatorname{mod} z^{N+1}$ で計算すれば十分です (~~さらに細かくいうと $\hat{k}(A)$ は $\operatorname{mod} z^{\lvert A\rvert+1}$ で合っていればよいです~~ ※追記これ普通に嘘です例を見てください)．

一応例もやっておきます． $N = 2$ として， $\mathcal{P}$ 上の関数は (入力形式のように) $\emptyset, \{0\}, \{1\}, \{0, 1\}$ の行き先の列で書くことにします．

$g = (1, 2, 3, 4),\, h = (5, 6, 7, 8)$ とする
集合のサイズの情報を持たせると $(1, 2z, 3z, 4z^2), (5, 6z, 7z, 8z^2)$
累積和をとって (いわゆるゼータ変換) $\hat{g} = (1, 1+2z, 1+3z, 1+5z+4z^2),\, \hat{h} = (5, 5+6z, 5+7z, 5+13z+8z^2)$
各点積をとって $\hat{k} = (5, 5+16z+12z^2, 5+22z+21z^2, 5+38z+93z^2+92z^3+32z^4)$
差分をとって (Möbius 変換) $k = (5, 16z+12z^2, 22z+21z^2, 60z^2+92z^3+32z^4)$
サイズが合っているところをとる $g h = (5, 16, 22, 60)$

さて，subset convolution は $2$ 個以外の積の場合も同様です (各点積をとるときに個数が変わるだけです)．また $\hat{}$ をつけたりとったりする操作は線型性を持っています．つまり例えば $5 g + 8 g^2 h$ を求めたければ $\hat{g}, \hat{h}$ を求めてから $5 \hat{g}(A) + 8 \hat{g}(A)^2 \hat{h}(A)$ を計算して変換を戻せばよいです．このように，多項式 $\phi$ に対し， $\phi(g)$ を計算するには各多項式 $\phi(\hat{g}(A))$ を求めればよいことになります．この多項式計算がそれぞれ $O(N^2)$ 時間でできれば全体で $O(2^N N^2)$ 時間となります．

sqrt

元の問題に戻りましょう． $f(A) = L(A) + \frac12 \sum_{B\sqcup C=A} f(B) f(C)$ は $f = L + \frac12 f^2$ と書けます．これを $f = 1 \pm \sqrt{1 - 2 L}$ のように解けたらよさそうなので，正当化します．

$L(\emptyset) = 0$ なので，subset convolution の定義を考えて $L^{N+1} = 0$ です．冪級数展開 $\sqrt{1 + x} = \sum_{i\ge0} \binom{1/2}{i} x^i = 1 + \frac{1}{2} x - \frac{1}{8} x^2 + \cdots$ に倣って $g = \sum_{i=0}^N \binom{1/2}{i} (-2L)^i$ と定めると， $g^2$ を計算したとき高次の項が消えて， $g^2 = 1 - 2 L$ がわかります． $g^2 = 1 - 2 L$ なる $g$ がこれかこれの $-1$ 倍しかないことは小さい部分集合から順番に値が決まっていくことから示せます．よって $g(\emptyset) = 1$ なる方を $\sqrt{1 - 2 L}$ と書けて， $f(\emptyset) = 0$ と合わせて $f = 1 - \sqrt{1 - 2 L}$ と求めたいものが書けました．

さて，あとは多項式 $\hat{L}(A)$ に対して $\sum_{i=0}^N \binom{1/2}{i} (-2\hat{L}(A))^i$ を計算できればよくなりました．これは形式的冪級数としての $\sqrt{1 - 2 \hat{L}(A)}$ を $\operatorname{mod} z^{N+1}$ で計算できればよく，それは低次の項から順番に合わせれば $O(N^2)$ 時間でできます．なお，説明の簡単のために値を $\mathbb{C}$ で考えましたが，今回は $2$ でしか割っていません (sqrt の代わりに exp 等が登場すると $1, \ldots, N$ の逆数が要ります)．

実装

$\Theta(3^N)$ を落とさなければいけない都合上 $O(2^N N^2)$ でも定数倍が幾分厳しめな時間制限設定になってしまっています．準備陣想定解のうち速いものは 1.5 秒程度 (C++, D)， $\Theta(3^N)$ のものは 12 秒程度 (C++) でした．想定解の場合は以下のような注意点が挙げられます：

高々 $N$ 次の多項式を $2^N$ 個持つので 2 次元配列を使うことになりますが，sqrt の計算が支配的なので，C++ などでいう [2^N][N+1] の順番のほうが [N+1][2^N] の順番より速くなるはずです．
形式的冪級数の sqrt は $O(N \log N)$ 時間でできますが，今回は $N$ が小さいので定数倍が非常に痛く向いていません．
sqrt をとるとき log をとって $\frac12$ 倍して exp するのは定数倍が大きく間に合わないかもしれません．log を介さず累乗を $O(N^2)$ で行う方法 (微分の利用) だとちょっとだけ遅いですが間に合います．
sqrt を低次の項から合わせるとき対称性を使って掛け算の回数を約半分にできます．これは必須ではありませんが結構効いてきます．
配列を静的に確保すると速くなります．なお，Java の 2 次元配列で動的に確保して 6 秒にぎりぎり間に合うくらいでした．

別解

heno239 さんのアイデアが，subset convolution を中身をいじらずに用いる方法でした．

dp[i]:=根がiになる、条件を満たす木の個数(各iについて葉としてふさわしいならdp[i]=1,そうでないならdp[i]=0で初期化)
dp2[i]:=Σ_{s in i}dp[s]*dp2[s^i](ただしdp2[0]=1で初期化)
として、各k(0<=k<N)について、
— heno (@heno_code) 2020年12月24日

リンク先 Twitter を見ていただきたいのですが，dp2 として「条件を満たす木において葉の $1$ 個に印をつけたもの」を数えています．これを一緒に計算することで， $\{0, \ldots, N-1\}$ の元を $1$ 個特別にとって再帰的に解くことができています．畳み込み方としてはいわゆる分割統治 FFT の気持ちにちょっと似ているかもしれません．

代数的には一体何が起こっているのかというのを補足しておきます． $t\colon \mathcal{P} \to \mathbb{C}$ を $t(\{N-1\}) = 1$ ， $t(A) = 0$ ( $A \ne \{N-1\}$ ) で定めて $L = L_0 + L_1 t,\, f = f_0 + f_1 t$ と書けます ( $N-1 \in A$ のとき $L_0(A) = L_1(A) = f_0(A) = f_1(A) = 0$ )． $f = L + \frac12 f^2$ を展開して， $t^2 = 0$ に注意すると， $f_0 = L_0 + \frac12 f_0^2$ および $f_1 = \frac{L_1}{1 - f_0}$ が得られます．前者は DP 配列の前半が再帰的に計算できることを言っています．後者は DP 配列の後半を求めるために $\frac{1}{1 - f}$ もほしいということを言っています． $g = \frac{1}{1-f}$ とおいて (これが dp2 です) 同様に $g = g_0 + g_1 t$ と書くと， $(1 - f) g = 1$ から $g_0 = \frac{1}{1 - f_0}$ および $g_1 = \frac{f_1 g_0}{1 - f_0} = f_1 g_0^2$ が得られます．これで $f_1, g_1$ を $L_1, g_0$ から求めることができました．一般に $\phi(f) = 0$ を解くときは $\phi(f_0) + \phi'(f_0) f_1 t = 0$ となり，これはつまり Newton 法の一種です．

subset convolution を中身をいじらず，とは言いましたが， $O(2^N N^2)$ 時間ではあるもののちょっと定数倍が重く残念ながら TLE になってしまったようです．いろいろ試してみたのですが，同じ列に対するゼータ変換・Möbius 変換を複数回行わないようにすることでなんとか通せるかもしれないというくらいでした．ちなみに，この解法には $2$ による割り算すら登場しません．

統計情報

正解者：チーム iwiwi (0:49:32) はじめ 4 チーム

高度典型です！

特に中国で元々有名だった手法です．あまりちゃんと調査を行っていないのですが，中国の IOI 選抜に使われる論文 [2] がこの手の応用のおそらくきっかけです．中国語で「子集卷积」 (subset convolution) やら「集合幂级数」 (set power series) やらを検索することで情報がわんさか出てきますね．日本で話題になったのは ARC 105 のときで，Elegia さんの書き込みのおかげで知れ渡る形となりました．

問題としては解法ドリブンみたいな作問ですが，そこそこ自然な設定で解けた・sqrt は例題が全然見当たらなかったので，日頃からアンテナを張ってる方々へのボーナスという気持ちも込めて出題しました．普段のコンテストだとそもそも計算量の区別でも不幸を呼びやすく出しにくいですし，クリスマスということで．

参考文献

[1] Björklund, Andreas, et al. "Fourier meets Möbius: fast subset convolution." Proceedings of the thirty-ninth annual ACM symposium on Theory of computing. 2007.

[2] 吕凯风. "集合幂级数的性质与应用及其快速算法." 2015年信息学奥林匹克中国国家队候选队员论文. 2015.

hos.lyric's blog