276. 栅栏涂色

276. 栅栏涂色

有 k 种颜色的涂料和一个包含 n 个栅栏柱的栅栏，请你按下述规则为栅栏设计涂色方案：

• 每个栅栏柱可以用其中 一种 颜色进行上色。

• 相邻的栅栏柱 最多连续两个 颜色相同。

给你两个整数 k 和 n ，返回所有有效的涂色 方案数 。

输入：n = 3, k = 2
输出：6
解释：所有的可能涂色方案如上图所示。注意，全涂红或者全涂绿的方案属于无效方案，因为相邻的栅栏柱 最多连续两个 颜色相同。

解法一：标准 DP

思想

状态表示

1. f[i][0] 表示第 i 个与第 i - 1 个栅栏同色
2. f[i][1] 表示第 i 个与第 i - 1 个栅栏异色

状态转移

1. 第 i 个与第 i - 1 个同色。只能从第 i - 1 个与第 i - 2 个不同色转移过来，且只能选择与上次相同的颜色（只有 1 种选择方式）。即：f[i][0] = f[i - 1][1]
2. 第 i 个与第 i - 1 个不同色。可以从第 i - 1 个与第 i - 2 个同色或异色转移，且选择的颜色只要与上次不同的就行（有 k - 1 种选择方式）。即：f[i][1] = (f[i - 1][0] + f[i - 1][1]) * (k - 1)

复杂度

时间复杂度：$O(N)$

空间复杂度：$O(N)$

代码

class Solution {
    public int numWays(int n, int k) {        
        if (k == 1) return n > 2 ? 0 : 1;
        
        int[][] f = new int[n][2];
        f[0][0] = 0;
        f[0][1] = k;
        for (int i = 1; i < n; i++) {
            f[i][0] = f[i - 1][1];
            f[i][1] = (f[i - 1][0] + f[i - 1][1]) * (k - 1);
        }

        return f[n - 1][0] + f[n - 1][1];
    }
}

解法二：滚动变量优化

思想

f[i][0] = f[i - 1][1];
f[i][1] = (f[i - 1][0] + f[i - 1][1]) * (k - 1);

发现只与 i - 1 层有关，于是可以用两组 2 个变量滚动更新

复杂度

时间复杂度：$O(N)$

空间复杂度：$O(1)$

代码

class Solution {
    public int numWays(int n, int k) {        
        if (k == 1) return n > 2 ? 0 : 1;

        int a = 0, b = k;
        for (int i = 1; i < n; i++) {
            int x = a, y = b;
            a = y;
            b = (x + y) * (k - 1);
        }
        
        return a + b;
    }
}