chineristACの日記

yukicoder No.1933 ABCstring 別解

おもしろかったです。writer/testerの方法より多少思いつきやすそうな&文字の種類増やせる別解で解いたので書きます（というか解説の固定方法よく思いついたな…）。

問題の一般化

以下のような問題が解ければいいです。

$N$ 種類の色 $i$ があり、色 $i$ のボールは $A_i$ 個、色 $i$ の筒は $B_i$ 個あります。同じ色の筒には同じ色のボールを入れることができませんが、すべての色のボールが入れられる特別な筒が $1$ つだけあります。ボールの入れ方が何通りあるか求めてください。

解法

$S=\sum A_i,\ T=\sum B_i,\ U=S+T$ とします。簡単のためにまず $1 \leq B_i$ とします。

ボールは同じ色であってもすべて区別します（最後に $\prod_{i=1}^{N}A_i!$ で割ればいいです）。そして「同じ色のボールは同じ色の筒に入れてはならない」という条件について、条件違反をボールごとにみる包除原理で考えます。

色 $i$ のボールのうち、確実に条件違反しているボールが $a_i$ 個であるものとします。ボールの筒への入れ方ですが、確実に条件違反しているボール→残りのボールの順に考えます。

まず確実に条件違反しているボールの入れ方ですがこれは色ごとに独立に考えることができて $\prod_{i=1}^{N}\frac{(a_i+B_i-1)!}{(B_i-1)!}$ 通りです。

この後残っているボールを挿入しますが、これは $\sum a_i + T$ 個のものの列に $S - \sum a_i$ 個のものを挿入と考えれば $\frac{(S+T)!}{(S-\sum a_i)!}$ 通りです。

よって $\sum a_i$ ごとに最初の総積の総和が求まればいいです。これは包除原理による符号や最初に違反する $a_i$ 個の選び方も考慮すると $f_i(x)=\sum _ {k=0}^ {A _ i} \binom{A _ i}{k}\frac{(k+B _ i-1)!}{(B _ i-1)!} (-x)^{k}$ 　の総積をみればよく、 $O(U\log^ 2 U)$ で求まります。

$B_i=0$ の $i$ については違反しようがないので $f_i(x)=1$ として同じ計算をすれば求まります。