問題
解法
数列は 0-indexed とする。
$\displaystyle \mathcal{B} _ {i, p} = \left\lbrace (B _ 0, \ldots, B _ {i - 1}) \in \lbrace 0, 1\rbrace ^ i \mid B _ j + B _ {j + 1} \neq 2 \ (\forall j\in\lbrace 0,1,\ldots, i-2\rbrace) \land B _ {i - 1} = p \right\rbrace$ とする。これは、ラフに言えば、前 $i$ 個のワインに対する盗む(=1)/盗まない(=0)の valid な決め方であって、ワイン $i - 1$ を盗むかどうかが $p$ であるようなもの全体の集合である。
以下のような、多項式を要素に持つ動的計画法を考える。ただし組 $(i, p, t)$ は $i \in \lbrace 0, 1, \ldots, n\rbrace, p\in\lbrace 0, 1\rbrace, t\in\lbrace 0, 1\rbrace$ の範囲を動く。
$$ f _ {i, p, t}(x) \coloneqq \sum _ {B \in \mathcal{B} _ {i, p}} \sum _ {j = 0} ^ {i - 1} B _ j A _ j ^ t x ^ {\sum _ {j = 0} ^ {i - 1} B _ j}. $$
$x$ の指数 $\sum _ {j = 0} ^ {i - 1} B _ j$ は盗むワインの個数に対応する。$p\in \lbrace 0, 1\rbrace$ を状態に持つのは次のワインを盗めるかの判定のためである。$t\in\lbrace 0, 1\rbrace$ は $A$ の $0$ 乗和と $1$ 乗和をそれぞれ持つことを意味する。
求める答えは $\lbrack x ^ K \rbrack (f _ {n, 0, 1} + f _ {n, 1, 1})$ である。遷移は次のように行列で表現できる。
$$\left(\begin{array}{cc:cc} 1 & 0 & 1 & 0 \\ 0 & 1 & 0 & 1 \\ \hdashline x & 0 & 0 & 0 \\ A _ ix & x & 0 & 0 \\ \end{array}\right)\begin{pmatrix} f _ {i, 0, 0} \\ f _ {i, 0, 1} \\ \hdashline f _ {i, 1, 0} \\ f _ {i, 1, 1} \end{pmatrix} = \begin{pmatrix} f _ {i+1, 0, 0} \\ f _ {i+1, 0, 1} \\ \hdashline f _ {i+1, 1, 0} \\ f _ {i+1, 1, 1} \end{pmatrix}.$$
上 $2$ 行が $i$ 番目のワインを盗まない場合の遷移、下 $2$ 行が $i$ 番目のワインを盗む場合の遷移に対応する。左 $2$ 列は $i - 1$ 番目のワインを盗まない場合からの遷移、右 $2$ 列は $i - 1$ 番目のワインを盗む場合からの遷移を表す。
従って $T _ i \coloneqq \begin{pmatrix} 1 & 0 & 1 & 0 \\ 0 & 1 & 0 & 1 \\ x & 0 & 0 & 0 \\ A _ ix & x & 0 & 0 \\ \end{pmatrix}$ と定めれば次が成り立つ。
$$T _ {n - 1} T _ {n - 2} \cdots T _ 0\begin{pmatrix}1 \\ 0 \\ 0 \\ 0\end{pmatrix} = \begin{pmatrix} f _ {n, 0, 0} \\ f _ {n, 0, 1} \\ f _ {n, 1, 0} \\ f _ {n, 1, 1} \end{pmatrix}.$$
$T _ 0, T _ 1, \ldots, T _ {n - 1}$ を順にベクトルに掛けると $\Theta(n ^ 2)$ 時間掛かるが、二分木状に積を取る分割統治法を用いて $T _ {n - 1} T _ {n - 2} \cdots T _ 0$ を計算してからベクトルに掛けることで (定数倍の非常に重い) $\Theta(n (\log n) ^ 2)$ 時間で本問題を解くことができる。
なお、本解法では全ての $K = 0, 1, \ldots, \lceil N/2\rceil$ に対する答えを計算できる。
実装
https://atcoder.jp/contests/arc120/submissions/42033491
実行時間制限に間に合わせるため、いくつか定数倍高速化を行っている。
- $T _ i$ の積として得られる行列は全て $\begin{pmatrix} \ast & 0 & \ast & 0 \\ \ast & \ast & \ast & \ast \\ \ast & 0 & \ast & 0 \\ \ast & \ast & \ast & \ast \\ \end{pmatrix}$ という形を持つため、$0$ であることが予め分かっている要素の計算を省いてよい
- 分割統治の過程において次数が $2$ 冪の多項式が多く現れる。$2 ^ k$ 次の多項式どうしの積は $2 ^ {k + 1}$ 次であり、配列のサイズとしては $2 ^ {k + 1} + 1$ となるため FFT の定数倍が悪化する。そこで一方の $2 ^ k$ 次の項を削って畳み込んだ後に削った項の寄与を足しこむことで、定数倍を削減できる。
