D - Matrix Pow Sum Editorial
by 
maspy
ヒント → https://atcoder.jp/contests/arc190/editorial/11715
[1] 方針
行列の \(0\) をそれぞれ適当な不定元に置き換えたうえで \(p\) 乗することを考えます．例えば入力例 1 であれば
\[\begin{pmatrix}x_{11}&1\\\ x_{21}&2\end{pmatrix}^3=\begin{pmatrix}x_{11}^3+2x_{11}x_{21}+2x_{21}&x_{11}^2+2x_{11}+x_{21}+4\\\ x_{11}^2x_{21}+x_{21}^2+2x_{11}x_{21}+4x_{21}&x_{11}x_{21}+4x_{21}+8\end{pmatrix}\]
のようになります．不定元に対して \(1, 2, \ldots, p-1\) に置き換えて和をとることは，計算結果の行列において不定元のべき乗 \(x_{ij}^k\) を \(1^k+\cdots+(p-1)^k\) に置き換えることに相当します．
ここで次が成り立ちます：
\(1^k+\cdots+(p-1)^k \equiv \begin{cases}-1 & (p-1\mid k) \\\ 0 & (p-1\nmid k)\end{cases} \pmod{p}\)
これは \(\mod p\) での原始根 \(r\) をとると，左辺が \(\sum_{i=0}^{p-2}r^{ki}\) に等しいことから証明できます．
したがって，どの不定元についての指数も \(p-1\) の倍数であるような項の係数をすべて求めて，その \((-1)^K\) 倍を答に加算すればよいです．
[2] 答の計算
\(p=2\) のときは答は定義通りに簡単に計算できます．以下 \(p\geq 3\) を仮定します．
次の \(2\) つのタイプの項があります．
- どの不定元についての指数も \(0\) である．
- ある不定元 \(x_{ij}\) についての指数が \(p-1\)，それ以外の不定元についての指数は \(0\) である．
前者のタイプについては入力の行列をそのまま \(p\) 乗することで計算できます．
後者のタイプについて，不定元が対角成分である場合に
- \(A_{ji}\cdot x_{ii}\cdots x_{ii}\)
- \(x_{ii}\cdots x_{ii}\cdot A_{ij}\)
という項があります．\(p\geq 5\) の場合には不定元が非対角成分である場合の項はありません．\(p=3\) の場合にはさらに
- \(x_{ij}\cdot A_{ji}\cdot x_{ij}\)
という項があります．これらをすべて計算することで，計算量 \(O(N^3\log p)\) で本問題の答が計算できます．
posted:
last update: