LintCode 947: Matrix Power Series (矩阵快速幂经典题)

947. Matrix Power Series
     中文English
     Given a n × n matrix A and a positive integer k and a positive integer m, find the sum S = A + A^2 + A^3 + … + A^k,return the elements of S modulo m

Example
Example 1:

Input:
[[0,1],[1,1]]
2
4
Output:
[[1,2],[3,4]]
Example 2:

Input:
[[1,1],[1,1]]
7
10
Output:
[[7,7],[7,7]]
Notice
n<=30,1<=k<=10^9,1<=m<104
the range of the element of the Matrix A is [0,10000]

解法1：矩阵快速幂
当k是偶数时：

S(k) = A + A ^ 2 + ... + A ^ k
     = A + A ^ 2 + A ^ (k/2) + A ^ (k/2 + 1) + A ^ (k/2 + 2) + A ^ k
     = (E + A ^ (k/2)) *  (A + A ^ 2 + ... + A ^ (k/2))
     = (E + A ^ (k/2)) * S(k/2)

当k是奇数时，我们有s(k) = s(k - 1) + A ^ k
注意：

1. S(k) = (E + A ^ (k/2)) * S(k/2)，其中的+号两边和乘号两边都可互换。
   因为当矩阵A * B的时候，如果A和B中有一个是单元矩阵，或A=B时满足交换律。
2. operator * 里面的最内层循环，可以先加完之后再模，这样省时间。当然前提是保证不溢出。

代码如下：

class Matrix {
public:
    vector<vector<long long>> matrix;

    Matrix(int m, int n, int k) {
        this->m = m;
        this->n = n;
        this->k = k;
        matrix.resize(n, vector<long long>(n, 0LL));
    }
    
    void unit() {
          for (int i = 0; i < n; ++i)
            for (int j = 0; j < n; ++j)
              if (i == j) matrix[i][j] = 1;
    }

    Matrix operator * (const Matrix & M) const {
        Matrix result(m, n, k);
        for (int i = 0; i < n; ++i) {
            for (int j = 0; j < n; ++j) {
                for (int h = 0; h < n; ++h) {
                    result.matrix[i][j] += matrix[i][h] * M.matrix[h][j];
                   // result.matrix[i][j] %= m;      
                }
                result.matrix[i][j] %= m;
            }
        }
        return result;
    }

    Matrix operator + (const Matrix & M) const {
        Matrix result(m, n, k);
        for (int i = 0; i < n; ++i) {
            for (int j = 0; j < n; ++j) {
                result.matrix[i][j] = (matrix[i][j] + M.matrix[i][j]) % m;
            }
        }
        return result;
    }
 
    Matrix power(int p) {
        Matrix result(m, n, k);
        result.unit();
        Matrix B = *this;
        while(p) {
            if (p & 0x1) result = result * B;
            B = B * B;
            p /= 2;
        }
        return result;
    }
    
    Matrix sum(int s) {
        if (s == 1) return *this;
        if (s & 0x1) {
            return sum(s - 1) + power(s);
        } else {
           Matrix E(m, n, k);
           E.unit();
           return (power(s / 2) + E) * sum(s / 2);
        }
    }
private:
    int m, n, k;
};

class Solution {
public:
    /**
     * @param A: a n x n matrix
     * @param k: a integer
     * @param m: a integer
     * @return: return a 2D arrays, denote S matrix
     */
    vector<vector<int>> matrixPowerSeries(vector<vector<int>> &A, int k, int m) {
        int n = A.size();
        vector<vector<int>> result;
        if (n == 0) return result;
        result.resize(n, vector<int>(n, 0));
        Matrix M(m, n, k);
        
        for (int i = 0; i < n; ++i) {
            for (int j = 0;  j < n; ++j) {
                M.matrix[i][j] = A[i][j];
            }
        }
        
        Matrix Mk = M.sum(k);

        for (int i = 0; i < n; ++i) {
            for (int j = 0; j < n; ++j) {
                result[i][j] = Mk.matrix[i][j];
            }
        }
        return result;
    }
};