lintcode python 代码 514 栅栏染色

我们有一个栅栏，它有n个柱子，现在要给柱子染色，有k种颜色可以染。
必须保证任意两个相邻的柱子颜色不同，求有多少种染色方案。
思路：第一个柱子有k中方案 number[1]=k
第二个柱子有k中方案 number[2]=k*k
第三个柱子： 1.和第二个柱子同色 和第一个不同色 number[2]=k*（k-1）
2.和第二个柱子不同色 number[2]=k*k*（k-1）
第n个柱子number[n]=record[n- 2] * (k - 1) + record[n-1] * (k - 1)

class Solution:
    # @param {int} n non-negative integer, n posts
    # @param {int} k non-negative integer, k colors
    # @return {int} an integer, the total number of ways
    def numWays(self, n, k):
        # Write your code here
        if n == 1:
            return k
        if n == 2:
            return k * k
        if k <= 1:
            return 0
        ans = [0] * n
        ans[0] = k
        ans[1] = k * k
        for i in range(2, n):
            ans[i] = ans[i-1] * (k-1) + ans[i-2] * (k-1)
        return ans[n-1]