Official

D - Matrix Pow Sum Editorial by evima


Hints → https://atcoder.jp/contests/arc190/editorial/11918


[1] Strategy

Consider replacing every \(0\) in the matrix by some indeterminate, and then taking the \(p\)-th power. For example, in Sample Input 1 we have:

\[\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}\]

and so on. Summing over all possibilities of substituting \(1, 2, \ldots, p-1\) for each indeterminate amounts to, in the resulting matrix, replacing each occurrence of \(x_{ij}^k\) by \(1^k + 2^k + \cdots + (p-1)^k\).

Here, the following holds: \(1^k+\cdots+(p-1)^k \equiv \begin{cases}-1 & (p-1\mid k) \\\ 0 & (p-1\nmid k)\end{cases} \pmod{p}.\)

This can be proved by taking a primitive root \(r\) modulo \(p\) and observing that the left hand side equals \(\sum_{i=0}^{p-2}r^{ki}\).

Hence, we just need to find the coefficients of those terms in which the exponent of every indeterminate is a multiple of \(p-1\), and add those terms multiplied by \((-1)^K\) to the answer.


[2] Computing the answer

When \(p=2\), the solution can be found straightforwardly by definition. Below we assume \(p \ge 3\).

There are two types of terms to consider:

  • Those for which all indeterminates have exponent \(0\).
  • Those for which exactly one indeterminate \(x_{ij}\) has exponent \(p-1\), and all the others have exponent \(0\).

For the first type, simply take the matrix \(A\) as given, raise it to \(p\), and compute that result.

For the second type, when the indeterminate is on the diagonal, we have the following terms:

  • \(A_{ji}\cdot x_{ii}\cdots x_{ii}\)
  • \(x_{ii}\cdots x_{ii}\cdot A_{ij}\)

If \(p \ge 5\), there are no terms corresponding to non-diagonal indeterminates. If \(p=3\), we also have the term:

  • \(x_{ij} \cdot A_{ji} \cdot x_{ij}\)

By enumerating all these, we can compute the answer in \(O(N^3 \log p)\) time.


posted:
last update: