【ACWing】225. 矩阵幂求和

题目地址：

https://www.acwing.com/problem/content/description/227/

给定$n\times n$矩阵$A$和正整数$k$，求和$S=A+A^2+A^3+\dots +A^k$。

输入格式：
输入只包含一个测试用例。
第一行输入包含三个正整数$n$，$k$和$m$。
接下来$n$行，每行包含$n$个非负整数（均不超过$32,768$），用以描绘矩阵$A$。

输出格式：
按与描述矩阵$A$相同的方式，输出将$S$中所有元素对$m$取模后得到的矩阵。

数据范围：
$1\le n\le 30$
$1\le k\le 10^9$
$1\le m<10^4$

设$I$为$n$阶单位阵，构造$B=\left[\begin{array}{cc}A& I\\ 0& I\end{array}\right]$，则$$B^2=\left[\begin{array}{cc}A& I\\ 0& I\end{array}\right]\left[\begin{array}{cc}A& I\\ 0& I\end{array}\right]=\left[\begin{array}{cc}{A}^2& A+I\\ 0& I\end{array}\right]$$以此类推，可由数学归纳法知道$B^{k+1}$的右上角为$A^k+A^{k-1}+...+I$，将这个矩阵减去$I$即得到答案。求$B^k$可以用快速幂。代码如下：

#include <cstring>
#include <iostream>
using namespace std;

const int N = 65;
int n, k, MOD;
int a[N][N], res[N][N];

void mult(int a[][N], int b[][N]) {
  static int c[N][N];
  memset(c, 0, sizeof c);
  for (int i = 1; i <= n; i++)
    for (int j = 1; j <= n; j++)
      for (int k = 1; k <= n; k++)
        c[i][j] = ((c[i][j] + a[i][k] * b[k][j]) % MOD + MOD) % MOD;

  for (int i = 1; i <= n; i++)
    for (int j = 1; j <= n; j++) a[i][j] = c[i][j];
}

void fast_pow(int k) {
  while (k) {
    if (k & 1) mult(res, a);
    k >>= 1;
    mult(a, a);
  }
}

int main() {
  scanf("%d%d%d", &n, &k, &MOD);
  for (int i = 1; i <= 2 * n; i++)
    for (int j = 1; j <= 2 * n; j++) {
      if (i <= n && j <= n) scanf("%d", &a[i][j]);
      else if (j - i == n || i > n && i == j)
        a[i][j] = 1;

      if (i == j) res[i][i] = 1;
    }

  n <<= 1;
  fast_pow(k + 1);
  n >>= 1;
  for (int i = 1; i <= n; i++)
    res[i][i + n] = ((res[i][i + n] - 1) % MOD + MOD) % MOD;

  for (int i = 1; i <= n; i++) {
    for (int j = n + 1; j <= 2 * n; j++) printf("%d ", res[i][j]);
    puts("");
  }
}

时间复杂度$O(n^3\log k)$，空间$O(n^2)$。