[LintCode 514] 柵欄染色(Python)

題目描述

我們有一個柵欄，它有n個柱子，現在要給柱子染色，有k種顏色可以染。
必須保證不存在超過2個相鄰的柱子顏色相同，求有多少種染色方案。（n和k都是非負整數）

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

               post 1    post 2   post 3
way 1           0         0        1
way 2           0         1        0
way 3           0         1        1
way 4           1         0        0
way 5           1         0        1
way 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) 。