Codeforces Round #373 (Div. 1): C. Sasha and Array

解法はよくある感じなので、発展的内容として Kitamasa 法について書こうと思います。セグメント木に関しては書きません。

コンパニオン行列を用いてフィボナッチ数列は次のように計算できる。

$$ \begin{pmatrix} F_n \\ F_{n+1} \end{pmatrix}=\begin{pmatrix} 0 & 1 \\ 1 & 1 \end{pmatrix}^n \begin{pmatrix} F_0 \\ F_1 \end{pmatrix} $$

コンパニオン行列を $A$ としたとき、たとえば次のような式が成り立つ。

$$ F_1+F_2+F_2+F_3 = (A^1+A^2+A^2+A^3) (F_0 \; F_1)^{\mathrm{T}} の 0 番目の成分 $$

つまり、フィボナッチ数の和を求める代わりに、対応する行列の和を求めてあげても答えが求まる。

やりたいのは次の 2 種類の処理。

これらの操作を高速に処理するにはセグメント木を用いればよい。

さきほどは行列で話を進めていたが、多項式の話に落とすこともできる。いわゆる Kitamasa 法。

今回の行列は漸化式を見れば分かるように $A^2=A+E$ という関係式が成り立つので、$A^n=Q(A^2-A-E)+R=R$ のような変形ができる。つまり $A^n = A^n \;\mathrm{mod}\; (A^2-A-E)$ なので $A^n=a_0E+a_1A$ と表せる。

これは $A$ に関する任意の多項式は $a_0 E + a_1 A$ と表せるということを意味している。つまり $A^n$ や $A^5+3A^7$ とかはすべて 2 つの整数値の情報に圧縮できる。

$\phi(A),\psi(A)$ を $A$ に関する多項式としよう。このとき、その積は次のように計算できる。

$$\phi(A) \psi(A) = (a_0E+a_1A)(b_0E+b_1A) \; \mathrm{mod}\; (A^2-A-E)$$

このときやっているのは（1次式）×（1次式）の多項式乗算と多項式剰余。行列の係数にだけ興味があるので行列積は必要ない。

多項式乗算は簡単だけど、多項式剰余ってそこまで書く機会がないから少しつらい。けれど、次の関係式を使えば結構簡単に求められる。

$$ A^n = A^{n-2}A^2=A^{n-2}(E+A) $$

このことを使えば次数を 1 ずつ下げられる。フィボナッチ数列だと分かりづらいので、トリボナッチ数列でやってみよう。トリボナッチ数列の場合 2 次式で表せるので、積は 4 次式である。

$$\phi(A) \psi(A) = a_0E+a_1A+a_2A^2+a_3A^3+a_4A^4$$

まず 4 次式から 3 次式に変形する。

$$ \begin{align} \phi(A) \psi(A) &= a_0E+a_1A+a_2A^2+a_3A^3+a_4A(E+A+A^2)\\ &= a_0'E+a_1'A+a_2'A^2+a_3'A^3 \end{align} $$

3 次式を 2 次式に変形する。

$$ \begin{align} \phi(A) \psi(A) &= a_0'E+a_1'A+a_2'A^2+a_3'(E+A+A^2)\\ &= a_0''E+a_1''A+a_2''A^2 \end{align} $$

これで無事 mod A³-A²-A-E の値が求まった。

というわけで、多項式の議論だけで n 乗とかは計算できるのであった。多項式乗算も多項式剰余も素直にやって O(n²) なので、行列を使うより圧倒的に早い。メモリも O(n²) から O(n) に落ちて嬉しい（今回だと大差ないけど）。

法が整数だろうと多項式だろうと大体同じ性質が成り立つので扱いも楽。

529 ms。ちなみに本番は行列で通してます。

コードをフィボナッチに特化させればもう少し速くなるかもしれないけど、まあ満足。

セグメント木とコンパニオン行列を知ってれば簡単。知らないときついかも。