[LintCode 514] 栅栏染色(Python)

题目描述

我们有一个栅栏，它有n个柱子，现在要给柱子染色，有k种颜色可以染。
必须保证不存在超过2个相邻的柱子颜色相同，求有多少种染色方案。（n和k都是非负整数）

样例
n = 3, k = 2, return 6

               post 1    post 2   post 3way 1           0         0        1way 2           0         1        0way 3           0         1        1way 4           1         0        0way 5           1         0        1way 6           1         1        0

思路

动态规划问题。设f(i)为共i个柱子的方案数。对于第i个柱子，它有2种可能：和第i−1个柱子颜色相同，和第i−1个柱子颜色不同。则共i个柱子的方案数为这2种可能对应的方案数的和。

• 和第i−1个柱子颜色相同：那么i,i−1的颜色只要和i−2不同就可以了，共有k−1个颜色，对应的方案数为(k−1)∗f(i−2)
• 和第i−1个柱子颜色不同：那么i的颜色只要和i−1不同就可以了，共有k−1个颜色，对应的方案数为(k−1)∗f(i−1)

综上f(i)=(k−1)∗f(i−2)+(k−1)∗f(i−1)

初始条件：f(1)=k,f(2)=k2

代码

class Solution:    """    @param: n: non-negative integer, n posts    @param: k: non-negative integer, k colors    @return: an integer, the total number of ways    """    def numWays(self, n, k):        # write your code here        if n == 0:            return 0        # 初始条件        if n <= 2:            return k ** n        # 在栅栏数多于2个时，如果只有1个颜色，则方案数为0。需要注意这个情况的处理。        if k == 1:            return 0        p1 = k        p2 = k * k        for i in range(3, n + 1):            tmp = p2            p2 = p2 * (k - 1) + p1 * (k - 1)            p1 = tmp        return p2

复杂度分析

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