G - At Most 2 Colors Editorial by PCTprobability


\(C=1\) の場合は全てを同じ色で塗るしかないため、以降 \(C \ge 2\) を仮定します。

マス \(i\) を塗る色を \(A_i\) と置きます。ここで、\(A\) をランレングス圧縮して長さの情報だけを残した数列を \(B\) と置きます。

例えば、\(A=(1,1,1,2,2,3,2,2)\) であれば \(B=(3,2,1,2)\) となります。

\(B\) を固定したとき、条件を満たす \(A\) の個数を求めます。\(|B|=1\) の場合は \(C\) 通りなので、\(|B| \ge 2\) を仮定します。

\(B_1\) に使う色と \(B_2\) を使う色の組は \(B\) によらず \(C(C-1)\) 通りです。

その後、\(i \ge 3\) に対しては、\(B_{i-1} \le K-2\) であれば \(1\) 通り、そうでなければ \(C-1\) 通りがあります。

よって、まとめると \(C(C-1) \times (C-1)^{k}\) 通りとなります。ただし、\(k\)\(2 \le i \le N-1\) を満たす整数 \(i\) のうち、\(B_i \ge K-1\) を満たすものの個数ととします。

\(|B|=l\) とした場合の答えは形式的べき級数を用いると、\(C \times (C-1) \times \frac{x}{1-x} \times \left(\frac{x+(C-2)x^{K-1}}{1-x}\right)^{l-2} \times \frac{x}{1-x}\)\(x^N\) の係数となります。

よって、\(f(x) = \frac{x+(C-2)x^{K-1}}{1-x}\) と置くと、\(\frac{x^2}{(1-x)^2} \times \frac{1}{1-f(x)}\)\(x^N\) の係数を求められればよいです。これは整理すると、\(\frac{x^2}{(1-x)(1-2x-(C-2)x^{K-1})}\) となります。

あとはこの形式的べき級数の \(x^N\) の係数が得られればよいです。形式的べき級数の逆元を愚直に求めれば \(\mathrm{O}(N\log N)\)、この形式的べき級数の係数が疎であることを利用すると \(\mathrm{O}(N)\)、Bostan-Mori 法を使えば \(\mathrm{O}(K\log K \log N)\) でこの問題を解くことができます。

posted:
last update: