部分点解法

色ごとに状態を割り当てて \(1\) 回の操作ごとに状態を遷移することを考えると、答えは

• 対角成分は \(0\) でそれ以外の \(ij\) 成分は \(A_i\) である \(N \times N\) 行列 \(M\) 、および
• 第 \(1\) 成分が \(1\) で他の成分が \(0\) である長さ \(N\) のベクトル

を用いて \(\frac{1}{A_1}(M^T v)_1\) と表すことが出来るとわかります。

行列 \(M\) およびベクトル \(v\) に対して \(a_n = (M^n v)_1\) と置いた時、数列 \((a_n)\) は線形漸化的な数列 (c-recursive とも)、すなわちある線形漸化式を用いて表せる数列になることが知られています。
また、最小な線形漸化式の次数は行列 \(M\) のサイズで抑えられることも知られています。さらに、

• 数列の一般項(第 \(0\) 項から \(2n-2\) 項まで)から \(n\) 次の線形漸化式を \(\mathrm{O}(n^2)\) で復元できる (Berlekamp-Massey 法) ことや、
• \(n\) 次の線形漸化的な数列の第 \(k\) 項を \(\mathrm{O}(n \log n \log k)\) で求められる (Bostan-Mori 法や Fiduccia 法など)

ことが知られています。(特に Berlekamp-Massey 法と Bostan-Mori 法を合わせて俗に BM-BM と呼ぶ人が多いです。詳細は maspy さんの記事 などを参照してください。)

よって今回の問題の場合、\((M^i v)_1\) を \(0 \leq i \leq 2N\) 程度の範囲で計算できれば上記のアルゴリズムで \(\mathrm{O}(N^2 + N \log N \log T)\) で答えを求められることがわかります。\((M^i v)_1\) の列挙は愚直に \(M\) を掛けていく方針だと \(\mathrm{O}(N^3)\) かかりますが、\(M\) が非常にキレイな行列である性質を生かすと計算量を \(\mathrm{O}(N^2)\) に落とすことができます。以上よりこの問題を \(\mathrm{O}(N^2 + N \log N \log T)\) で解くことが出来ます。

満点解法

この問題は隣接する要素が異なる \(N\) 以下の正整数からなる数列 \(B_0=1, B_1, \dots, B_T=1\) に対する \(\prod_{i=1}^{T-1} A_{B_i}\) の総和を計算する問題だと考えることが出来ます。

両端が \(i\) という条件は扱いづらいので片端が \(i\) という条件に変えましょう。

\(F_i(x)\) を「長さが \(n\) で最終項が \(i\) である数列 \(B_0, B_1, \dots, B_n=i\) に対する \(\prod_{i=1}^{T-1} A_{B_i}\) の総和」の母関数として定義します。すると \(F_i\) は

\[F_i(x) = A_i + A_i x \left(\sum_{1 \leq j \leq N, j \neq i} F_j \right)\]

と表せます。

\[G(x) = \sum_{i=1}^N F_i(x)\]

と置くと

\[F_i(x) = A_i + A_i x (G - F_i)\]

になり、変形して

\[F_i(x) = \frac{a_i(1 + xG(x))}{1 + a_i x}\]

となります。全ての \(i\) に対して上式の和を取って

\[G(x) = (1 + x G(x)) \left(\sum_{i=1}^N \frac{a_i}{1+a_i x} \right)\]

を得ます。

\[H(x) =\sum_{i=1}^N \frac{a_i}{1+a_i x} \]

と置いて上式を変形すると

\[G(x) = \frac{H(x)}{1 - x H(x)}\]

を得ます。\(H(x)\) は \(\mathrm{O}(N \log^2 N)\) で計算できるので \(G(x)\) および \(F_i(x)\) もまた \(\mathrm{O}(N \log^2 N)\) で計算できます。

\(F_1(x)\) が計算できると、少しの変形により問題の答えの母関数も計算できます。(これは読者への課題とします。) あとは Bostan-Mori などで答えを計算すればよく、 \(\mathrm{O}(N \log N (\log N + \log T))\) で問題を解くことが出来ます。