はじめに
ABC358のE問題を解いているときに母関数を使って解く方法を思いついたので書きたいと思います。
コンテスト中に通せませんでした（泣）
本題
問題文（前半のストーリーを省略）
長さ $1$ 以上 $K$ 以下の英大文字からなる文字列であって、以下の条件を満たすものの個数を $998244353$ で割った余りを求めてください。
- $1\le i\le 26$ を満たす任意の整数 $i$ について以下が成立する。
- 辞書順で $i$ 番目の英大文字を $a_i$ とおく。例えば、$a_1=$
A, $a_5=$E, $a_{26}=$Zである。 - 文字列の中に含まれている $a_i$ の個数は $0$ 個以上 $C_i$ 個以下である。
- 辞書順で $i$ 番目の英大文字を $a_i$ とおく。例えば、$a_1=$
制約
- $1\le K\le 1000$
- $0\le C_i\le 1000$
- 入力される値はすべて整数
解法
簡単に言えば、$26$ 種類のモノの重複順列であって、並べる長さとそれぞれのモノが含まれる個数に制限があるパターンを考えることになります。突然ですが、$f_i\ (i=1,2,\cdots ,26)$ と $F$ を以下のように定義してみましょう。 $$f_i(x)=\sum_{n=0}^{C_i}\dfrac{x^n}{n!},\quad F(x)=\prod_{i=1}^{26}f_i(x)$$ このとき、解答すべき数値は $\displaystyle\sum_{n=1}^{K}\ \lbrack x^n \rbrack F(x) \cdot n!$ を $998244353$ で割った余りであることがわかります。
説明
まずは具体的な状態で考えます。Aを$5$個、Bを$7$個、Cを$4$個並べるとします。この並べ方の総数は簡単な高校数学で求まりますね。$\dfrac{(5+7+4)!}{5!\ 7!\ 4!}$ 通りとなります。ここで分子（＝文字列の長さ）を $n$ と固定してみましょう。こうすると文字列の長さが $n$ のときの並べ方の総数は $0\le x_i\le C_i\ (i=1,2,\cdots ,26)$ かつ $x_1+x_2+\cdots +x_{26}=n$ を満たしながら $x_1,x_2,\cdots ,x_{26}$ を動かしたときの $\dfrac{n!}{x_1!x_2!\cdots x_{26}!}$ の総和となることがわかります。これは文字列長 $n$ を $x$ の指数で管理する母関数で表現できそうです。『 $\dfrac{1}{x_1!x_2!\cdots x_{26}!}$ の総和に $n!$ を掛ける』とするとうまくいきます。実際、$F(x)$ の $x^n$ の係数は $\dfrac{1}{x_1!x_2!\cdots x_{26}!}$ の総和となっています。
ACコード（Python）
- 多項式の畳み込みはナイーブなアルゴリズムでは $O(\sigma K^2)$ となります。（本問題では$K\le 1000$を満たすためこの方法でも間に合います。）このコードではより強い制約にも対応できるようにshakayami様のACL-for-pythonのconvolutionをお借りしてNTTというアルゴリズムで畳み込み計算を行っているため $O(\sigma K\log K)$ で計算が実行できています。（アルファベットの種類を $\sigma$ とおいています。）
- また、モジュラ逆元などの知識があるとよいです。
- このコードではおそらく $K=2\times 10^5$ などの場合でも間に合います。
def main():
MOD = 998244353
# 階乗計算
fac = [1]*(1000 + 9)
for a in range(1, 1000 + 4):
fac[a] = (fac[a-1]*a)%MOD
# 逆元計算
inv = [None, 1]
for a in range(2, 1000 + 9):
inv.append((-(MOD//a)*inv[MOD%a])%MOD)
# 階乗の逆元計算
facinv = [1]*(1000 + 9)
for a in range(2, 1000 + 4):
facinv[a] = (facinv[a-1]*inv[a])%MOD
fft = FFT(MOD)
K = int(input())
C = list(map(int,input().split()))
f = [1]
for i in range(26):
g = [facinv[n] for n in range(min(K+1, C[i]+1))]
f = fft.convolution(f, g)[:K+9]
ans = 0
for n in range(1, min(K+1, len(f))):
ans += f[n]*fac[n]
ans %= MOD
print(ans)
"""
ACL-for-python の FFT をコピペ
"""
if __name__ == "__main__":
main()
おわりに
- コンテスト時間中に考察を完了できる考察力を鍛えたいです。
- FFTと母関数を使えば本問題の並び替えでないバージョンなど、様々なバリエーションのある問題にも対応できそうだと思いました。