Paint Fence

题目描述

There is a fence with n posts, each post can be painted with one of the k colors.

You have to paint all the posts such that no more than two adjacent fence posts have the same color.

Return the total number of ways you can paint the fence.

解题方法

对于每一个fence,有两种选择

  • different from previous one, k - 1种选择
  • same as previous one, 1中选择

  • dp[i][j][0]表示第i个fence的颜色为j,同时与第i-1个fence颜色不一样的方案

  • dp[i][j][1]表示第i个fence的颜色为j,同时与第i-1个fence颜色一样的方案

因为只需要上一个i-1的value,所以可以简化成一个滚动数组。

Solution

简化版DP

直接两个变量表示之前的两种状态

class Solution(object):
    def numWays(self, n, k):
        """
        :type n: int
        :type k: int
        :rtype: int
        """
        if n == 0:
            return 0
        same = 0
        diff = k
        for i in range(2, n+1):
            t = diff
            diff = (same + diff) * (k-1) # previous total
            same = t

        return same + diff

Reference