【Leetcode】1411. Number of Ways to Paint N × 3 Grid

题目地址：

https://leetcode.com/problems/number-of-ways-to-paint-n-3-grid/

考虑对一个$n\times 3$的网格，每个格子可以涂$3$种颜色，要求相邻的两个格子不能涂同一种颜色。问涂色总共的方案数。答案模$10^9+7$返回。

以下用$0,1,2$三个数字来指代三种颜色。第一行有$12$种方案，其中只含两种颜色的方案有$6$种，分别是$010,020,101,121,202,212$，含$3$种颜色的方案也是$6$种，分别是$012,021,102,120,201,210$。设$f[k][0]$是第$k$行如果只涂两种颜色一共有多少合法方案，$f[k][1]$是第$k$行如果涂三种颜色一共有多少合法方案。考虑第$k-1$行，如果其涂两种颜色，例如$010$，那么第$k$行涂两种颜色的方案有$101,121,202$一共$3$种，第$k$行图三种颜色的方案有$102,201$一共$2$种；如果其涂三种颜色，例如$012$，那么第$k$行涂两种颜色的方案有$101,121$一共$2$种，第$k$行涂三种颜色的方案有$120,201$一共$2$种。注意到由于对称性，涂色方案只与涂了的颜色种数有关，与具体涂了什么颜色无关，所以$$\{\begin{array}{l}f[k][0]=3f[k-1][0]+2f[k-1][1]\\ f[k][1]=2f[k-1][0]+2f[k-1][1]\end{array}$$只需要递推一遍最后返回$f[n][0]+f[n][1]$即可。初始条件$f[1][0]=f[1][1]=6$。代码如下：

public class Solution {
    public int numOfWays(int n) {
        int MOD = (int) (1e9 + 7);
        long[][] f = new long[2][2];
        f[1][0] = f[1][1] = 6;
        for (int i = 2; i <= n; i++) {
            f[i & 1][0] = (f[i - 1 & 1][0] * 3 + f[i - 1 & 1][1] * 2) % MOD;
            f[i & 1][1] = (f[i - 1 & 1][0] * 2 + f[i - 1 & 1][1] * 2) % MOD;
        }
        
        return (int) ((f[n & 1][0] + f[n & 1][1]) % MOD);
    }
}

时间复杂度$O(n)$，空间$O(1)$。

可以用矩阵乘方 + 快速幂来做。若$A=\left[\begin{array}{cc}3& 2\\ 2& 2\end{array}\right]$，实际上就是要求$v=A^{n-1}\left[\begin{array}{c}6\\ 6\end{array}\right]$，最终返回$v_0+v_1$。算矩阵乘方可以用快速幂来做。代码如下：

public class Solution {
    
    class Matrix {
        private final int MOD = (int) (1e9 + 7);
        private long a11, a12, a21, a22;
        
        public Matrix(int a11, int a12, int a21, int a22) {
            this.a11 = a11;
            this.a12 = a12;
            this.a21 = a21;
            this.a22 = a22;
        }
        
        public void mult(Matrix B) {
            long c11 = (a11 * B.a11 + a12 * B.a21) % MOD;
            long c12 = (a11 * B.a12 + a12 * B.a22) % MOD;
            long c21 = (a21 * B.a11 + a22 * B.a21) % MOD;
            long c22 = (a21 * B.a12 + a22 * B.a22) % MOD;
            a11 = c11;
            a12 = c12;
            a21 = c21;
            a22 = c22;
        }
    }
    
    public int numOfWays(int n) {
        int MOD = (int) (1e9 + 7);
        Matrix res = new Matrix(1, 0, 0, 1), A = new Matrix(3, 2, 2, 2);
        n--;
        // 快速幂模板
        while (n > 0) {
            if ((n & 1) == 1) {
                res.mult(A);
            }
            
            n >>= 1;
            A.mult(A);
        }
        
        return (int) ((res.a11 + res.a12 + res.a21 + res.a22) * 6 % MOD);
    }
}

时间复杂度$O(\log n)$，空间$O(1)$。