HDU 4965 Fast Matrix Calculation 矩阵快速幂

题意：

给出一个\(n \times k\)的矩阵\(A\)和一个\(k \times n\)的矩阵\(B\)，其中\(4 \leq N \leq 1000, \, 2 \leq K \leq 6\)。

矩阵\(C=A \cdot B\)，求矩阵\(C^{N^2}\)的各个元素之和，以上矩阵运算均是在模\(6\)的情况下计算的。

分析：

如果我们直接计算\(A \cdot B\)的话，这个矩阵非常大，不可能进行快速幂计算。

所以要变形一下，

\((A \cdot B)^{N^2}=A \cdot (B \cdot A)^{N^2-1} \cdot B\)

而矩阵\(B \cdot A\)非常小，问题就解决了。

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

const int maxn = 1000 + 10;
const int MOD = 6;

inline int mul_mod(int a, int b) {
    return a * b % MOD;
}

inline void add_mod(int& a, int b) {
    a += b; if(a >= 6) a -= 6;
}

int N, K, A[maxn][MOD], B[MOD][maxn];
int tmp[maxn][MOD], res[maxn][maxn];

struct Matrix
{
    int a[MOD][MOD];

    Matrix() { memset(a, 0, sizeof(a)); }

    Matrix operator * (const Matrix& t) const {
        Matrix ans;
        for(int i = 0; i < K; i++)
            for(int j = 0; j < K; j++) if(a[i][j])
                for(int k = 0; k < K; k++)
                    add_mod(ans.a[i][k], mul_mod(a[i][j], t.a[j][k]));
        return ans;
    }
};

Matrix pow_mod(Matrix a, int n) {
    Matrix ans;
    for(int i = 0; i < K; i++) ans.a[i][i] = 1;
    while(n) {
        if(n & 1) ans = ans * a;
        a = a * a;
        n >>= 1;
    }
    return ans;
}

int main()
{
    while(scanf("%d%d", &N, &K) == 2 && N) {
        for(int i = 0; i < N; i++)
            for(int j = 0; j < K; j++)
                scanf("%d", &A[i][j]);
        for(int i = 0; i < K; i++)
            for(int j = 0; j < N; j++)
                scanf("%d", &B[i][j]);

        Matrix M;
        for(int i = 0; i < K; i++)
            for(int j = 0; j < N; j++) if(B[i][j])
                for(int k = 0; k < K; k++)
                    add_mod(M.a[i][k], mul_mod(B[i][j], A[j][k]));
        M = pow_mod(M, N * N - 1);

        memset(tmp, 0, sizeof(tmp));
        memset(res, 0, sizeof(res));
        for(int i = 0; i < N; i++)
            for(int j = 0; j < K; j++) if(A[i][j])
                for(int k = 0; k < K; k++)
                    add_mod(tmp[i][k], mul_mod(A[i][j], M.a[j][k]));
        for(int i = 0; i < N; i++)
            for(int j = 0; j < K; j++) if(tmp[i][j])
                for(int k = 0; k < N; k++)
                    add_mod(res[i][k], mul_mod(tmp[i][j], B[j][k]));

        int ans = 0;
        for(int i = 0; i < N; i++)
            for(int j = 0; j < N; j++)
                ans += res[i][j];
        printf("%d\n", ans);
    }

    return 0;
}