线性递推公式的矩阵快速幂技巧

快速幂

顾名思义， 快速幂是指快速求解幂运算的技巧。 正常求\(a^n\)的值需要执行n次相乘操作， 而快速幂能在\(log_2n\)时间复杂度内完成。

以求\(a^{27}\)为例， 27=1+2+8+16， 根据乘法结合律可得\(a^{27} = a^1*a^2*a^8*a^{16}\)， 即只需要指数转化为二进制并且求得对应位是1的幂再累计相乘就能快速得出结果。

快速幂的核心代码如下:

  /**
   * 求 x^n 的值且结果对mod取余
   */
  public static int power(int x, int n, int mod) {
      long res = 1, a = x;// 临时变量用long类型防止乘积超出整形范围
      while (n > 0) {
          if ((n & 1) == 1) {
              res = (res * a) % mod;
          }
          a = (a * a) % mod;
          n >>= 1;
      }
      return (int) res;
  }

矩阵乘法

按记忆里的内容简单说一下矩阵和矩阵乘法的概念， 有表述不严谨的地方请自行参考线性代数相关知识。

简单通俗来说， 矩阵就是m行n列的一组数字， 比如
\(A=\begin{pmatrix}1 & 2 & 3 \\ 4 & 5 & 6\end{pmatrix}\) 代表2行3列的矩阵。

当矩阵A的列数与矩阵B的行数相等时， A和B相乘才有意义。 乘积矩阵i行j列的值由矩阵A的i行与B的j列值依次相乘再累加得到， 即\((AB)_{i,j} = a_{i,1}b_{1,j}+a_{i,2}b_{2,j}+...+a_{i,n}b_{n,j}\)。

设\(A=\begin{pmatrix} a_1 & a_2 & a_3 \\ a_4 & a_5 & a_6 \end{pmatrix} \), \(B= \begin{pmatrix} b_1 & b_2 \\ b_3 & b_4 \\ b_5 & a_6 \end{pmatrix} \), 那么 \(A * B\) 的结果就是\( \begin{pmatrix} a_1b_1+a_2b_3+a_3b_5 & a_1b_2+a_2b_4+a_3b_6 \\ a_4b_1+a_5b_3+a_6b_5 & a_4b_2+a_5b_4+a_6b_6 \end{pmatrix} \)。

用二维数组来表示矩阵， 下面是矩阵相乘的核心代码：

  /**
   * 矩阵A和B的乘积, 结果对mod取余
   * A的列数必须等于B的行数
   */
  public static int[][] matrixMultiply(int[][] a, int[][] b, int mod) {
      int m = a.length, n = b[0].length;
      int k = b.length;
      int[][] res = new int[m][n];
      for (int aRow = 0; aRow < a.length; aRow++) {
          for (int bCol = 0; bCol < n; bCol++) {
              long num = 0;
              for (int i = 0; i < k; i++) {
                  num = (num + (long) a[aRow][i] * b[i][bCol]) % mod;
              }
              res[aRow][bCol] = (int) num;
          }
      }
      return res;
  }

矩阵快速幂

只有行列数相同的方形矩阵才能求幂， 矩阵的幂运算可以用相同技巧来快速求解。 这里再介绍一个单位矩阵 I 的概念， 简单说就是主对角线上值为1其它位置全部为0的方形矩阵， 任何矩阵与单位矩阵的乘积都不会改变， 等同于自然数中的1。

代码如下：

  /**
   * 求矩阵A的n次幂, 矩阵值对mod取余
   */
  public static int[][] matrixPower(int[][] a, int n, int mod) {
      int row = a.length;
      // res初始化为单位矩阵
      int[][] res = new int[row][row];
      for (int i = 0; i < row; i++) {
          res[i][i] = 1;
      }
      while (n > 0) {
          if ((n & 1) == 1) {
              res = matrixMultiply(a, res, mod);
          }
          a = matrixMultiply(a, a, mod);
          n >>= 1;
      }
      return res;
  }

应用场景

1. 固定关系线性递推的1维k阶问题快速求第i项， 比如斐波那契数、泰波那契数等；
2. 线性递推的k维1阶问题快速求第i项， 比如统计元音字母序列的数目、学生出勤记录II等；

相关题目

• 【LeetCode 509】斐波那契数列
• 【LeetCode 1137】第N个泰波那契数
• 【LeetCode 790】多米诺和托米诺平铺
• 【LeetCode 1220】统计元音字母序列的数目
• 【LeetCode 552】学生出勤记录II

参考链接：
左程云大佬的视频：算法讲解098【必备】快速幂、矩阵快速幂与两类问题