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: