LintCode 栅栏染色

题目描述：我们有一个栅栏，它有n个柱子，现在要给柱子染色，有k种颜色可以染。必须保证最多只有两个相邻的柱子颜色相同，求有多少种染色方案。
样例：n = 3, k = 2, return 6
post 1, post 2, post 3
way1 0 0 1
way2 0 1 0
way3 0 1 1
way4 1 0 0
way5 1 0 1
way6 1 1 0

动态规划求解：
用dp[i]保存染i个柱子的染色方案。dp[1]=k。dp[2]=k*k。
染第i个柱子时，和前面两个柱子相关.
前面2个柱子中，即第i-1,第i-2个柱子，设颜色相同的组有x组，颜色不同的组有dp[i-1]-x组。那么dp[i]=x*(k-1)+(dp[i-1]-x)*k.
如何计算第i-1，第i-2个柱子中，有多少颜色相同组，有如下：第1，2个柱子中不重复的等于第2,3个柱子中重复的组数。一般来说，就是第i-3，i-2个柱子中不重复的等于第i-1，第i-2个柱子中重复的组数。
第1，2个柱子中，重复的有k组，设x=k。
对于第3根及之后的柱子，计算dp[i]只需dp[i-1]，x。所以无需使用列表，只用2个变量即可。贴出代码如下：

def numWays(self, n, k):
        if n==0 or k==0:
            return 0
        elif n==1:
            return k
        elif k==1 and n>2:   #此句排除了只有一种颜色的可能
            return 0
        x=k        #第i-1，i-2个柱子中重复的组数
        a=k*k      #只用变量即可，a表示dp[i-1]，此处a初始化为dp[2]
        for i in range(3,n+1):
            b=(a-x)*k+x*(k-1)  #计算dp[i] 
            x=a-x  #更新x，用于下次计算
            a=b    #更新dp[i-1]，用于下次计算
        return a

这段代码在lintcode上提交的时候，有个问题，有一组数据，n=1367485（一个很大的数），k=1，提交之后说超出内存限制，然后为了排除这个问题，加上了上段代码中的下面这句。

elif k==1 and n>2:   #此句排除了只有一种颜色的可能
            return 0

还没搞清楚是什么原因。

另一种来自于别人的答案，就没有这种问题，看起来也比上一种好理解一些，见下面：
你看看别人家的孩子%>_<%