問題
「$N$ 個の頂点からなる順序付き根付き木」であって、各頂点の子の個数が $0$ 個以上 $K$ 個以下（子の並び順は区別する）であるものの個数を $998244353$ で割った余りを求めよ。
この数列の母関数 $y=y(x)$ は $y = x\cdot\phi(y)$、$\phi(y)=1+y+y^2+\cdots+y^K$ を満たす。ラグランジュの反転公式により
$$[x^N]y = \frac1N[y^{N-1}]\phi(y)^N$$
これを、多項式のべき乗を高速に計算する方法（NTTによる繰り返し二乗法）を用いて求めよ。愚直な畳み込みDPの $O(N^2)$ ではなく $O(N\log^2 N)$ 程度で計算せよ。
入力形式
N K制約
$1 \le N \le 10000$
$1 \le K \le 100$
法は $998244353$（NTT-friendly）
入出力例
入力例1
4 2出力例1
4入力例2
5 1出力例2
1概念図
ヒント（段階的開示）
ヒント1（方向性）
直接畳み込みで漸化式を回すと$O(N^2)$。求めたいのが特定の1つの$N$の値だけなら、母関数の構造を利用して途中経過を全部計算せずに直接その係数だけを式で表せないだろうか。
ヒント2（アプローチ）
「$y=x\cdot\phi(y)$」という陰関数の係数を直接取り出す道具がラグランジュの反転公式。$[x^N]y=\frac1N[y^{N-1}]\phi(y)^N$に変形すると、問題は「$\phi(y)^N$の$y^{N-1}$の係数」を求めることに帰着する。これは打ち切り付き多項式べき乗であり、繰り返し二乗法で$O(\log N)$回の乗算に帰着できる。各乗算をNTTで$O(N\log N)$処理すれば全体$O(N\log^2 N)$。
ヒント3（誘導）
def poly_pow_trunc(base, exponent, limit):
result = [1] + [0] * (limit - 1)
cur = base[:limit] + [0] * (limit - len(base))
e = exponent
while e > 0:
if e & 1:
result = poly_mult_trunc(result, cur, limit)
cur = poly_mult_trunc(cur, cur, limit)
e >>= 1
return result
phi = [1] * (K + 1)
q = poly_pow_trunc(phi, N, N)
answer = q[N - 1] * pow(N, MOD - 2, MOD) % MOD模範解答 (Python)
import sys
MOD = 998244353
G = 3
def ntt(a, invert):
n = len(a)
j = 0
for i in range(1, n):
bit = n >> 1
while j & bit:
j ^= bit
bit >>= 1
j ^= bit
if i < j:
a[i], a[j] = a[j], a[i]
length = 2
while length <= n:
w = pow(G, (MOD - 1) // length, MOD)
if invert:
w = pow(w, MOD - 2, MOD)
half = length >> 1
for i in range(0, n, length):
wn = 1
for k in range(half):
u = a[i + k]
v = a[i + k + half] * wn % MOD
a[i + k] = (u + v) % MOD
a[i + k + half] = (u - v) % MOD
wn = wn * w % MOD
length <<= 1
if invert:
n_inv = pow(n, MOD - 2, MOD)
for i in range(n):
a[i] = a[i] * n_inv % MOD
def poly_mult_trunc(a, b, limit):
result_len = len(a) + len(b) - 1
sz = 1
while sz < result_len:
sz <<= 1
fa = a + [0] * (sz - len(a))
fb = b + [0] * (sz - len(b))
ntt(fa, False)
ntt(fb, False)
for i in range(sz):
fa[i] = fa[i] * fb[i] % MOD
ntt(fa, True)
return fa[:limit]
def poly_pow_trunc(base, exponent, limit):
result = [1] + [0] * (limit - 1)
cur = base[:limit] + [0] * max(0, limit - len(base))
e = exponent
while e > 0:
if e & 1:
result = poly_mult_trunc(result, cur, limit)
cur = poly_mult_trunc(cur, cur, limit)
e >>= 1
return result
def main():
data = sys.stdin.buffer.read().split()
N = int(data[0])
K = int(data[1])
if N == 1:
print(1 % MOD)
return
phi = [1] * (K + 1)
limit = N
q = poly_pow_trunc(phi, N, limit)
coeff = q[N - 1]
inv_n = pow(N, MOD - 2, MOD)
print(coeff * inv_n % MOD)
main()Step-by-Step 解説
1自己相似構造から陰関数を立てる
「頂点数Nの木」は「根1個+0〜K個の部分木（同じ規則）」に分解できる。部分木の選び方の総和がφ(y)、根の分だけxを1つ掛けてy=x・φ(y)。
「頂点数Nの木」は「根1個+0〜K個の部分木（同じ規則）」に分解できる。部分木の選び方の総和がφ(y)、根の分だけxを1つ掛けてy=x・φ(y)。
2ラグランジュの反転公式の意味
陰関数の$x^N$係数を、yを求めずにφ(y)^Nという普通の多項式べき乗の係数として直接取り出せる。a(1)〜a(N-1)を経由せずa(N)単体を直接計算できる。
陰関数の$x^N$係数を、yを求めずにφ(y)^Nという普通の多項式べき乗の係数として直接取り出せる。a(1)〜a(N-1)を経由せずa(N)単体を直接計算できる。
3打ち切り付き多項式べき乗
y^{N-1}までの係数だけあればよい。繰り返し二乗法の各乗算結果を次数N-1で打ち切ることで、O(log N)回の乗算に抑える。
y^{N-1}までの係数だけあればよい。繰り返し二乗法の各乗算結果を次数N-1で打ち切ることで、O(log N)回の乗算に抑える。
4NTTによる多項式乗算
998244353は生成元3を持つNTT-friendly素数。長さが2べきの配列に対し畳み込みをO(n log n)で計算できる。
998244353は生成元3を持つNTT-friendly素数。長さが2べきの配列に対し畳み込みをO(n log n)で計算できる。
5最終的な係数の取り出しと逆元
φ(y)^Nのy^{N-1}係数に1/Nを掛けて答え。N=1は分母0の特殊ケースなので個別に1を出力する。
φ(y)^Nのy^{N-1}係数に1/Nを掛けて答え。N=1は分母0の特殊ケースなので個別に1を出力する。
よくあるミス
| ミス | 原因 | 正しい書き方 |
|---|---|---|
| N=1を公式にそのまま当てはめてゼロ除算 | 1/Nの分母が0になることを見落とす | N=1のときは個別に1を出力する |
| poly_pow_truncの打ち切り次数を誤る | y^{N-1}の係数だけ必要なのに配列長を余分に確保 | 配列長はN（インデックス0〜N-1）で十分 |
| NTTのサイズを2べきに切り上げ忘れる | 畳み込み結果の長さがちょうど2べきでない | poly_mult_trunc内で常にszを2べきに切り上げる |
| φ(y)の項数とlimitの大小関係を考慮しない | K+1>Nのときbase[:limit]で切り詰めが必要 | cur初期化で必ずlimitに合わせて切り詰め・パディング |
| 法を10**9+7にしてしまう | NTTに使えるのは998244353のような特殊な形の素数 | 本問では998244353を法として使う |
次のステップ
- 発展: φ(y)を一般の多項式（形式的べき級数）に差し替えれば任意の自己相似構造の数え上げに応用できる。K→∞（φ(y)=1/(1-y)）が通常のカタラン数系列に対応することを確認せよ。
- 発展2: 複数の異なるNに対して答えを求めたい場合、φ(y)^N=exp(N・log φ(y))というexp/log経由の定式化でどれだけ高速化できるか検討せよ。