276. Paint Fence

You are painting a fence of n posts with k different colors. You must paint the posts following these rules:

• Every post must be painted exactly one color.
• There cannot be three or more consecutive posts with the same color.

Given the two integers n and k, return the number of ways you can paint the fence.

class Solution {
public:
    int numWays(int n, int k) {
        if(n == 0) return n;
        if(n == 1) return k;
        int diff = k * (k-1);
        int eq = k;
        for(int i = 2; i < n; i++) {
            int tmp = diff;
            diff = (diff + eq) * (k-1);
            eq = tmp;
        }
        return diff + eq;
    }
};