Leetcode 276: Paint Fence (O(1) space)

leetcode_276c.cpp

#include <iostream>
#include <vector>

using namespace std;

int numWays(int n, int k) {
    if (n==0 || k==0) return 0;
    if (n==1) return k; 
    if (n==2) return k*k;
    vector<int> dp {0,1,k};

    for (int c=2; c<n; ++c) {
        dp[0]=dp[1]; dp[1]=dp[2];
        dp[2] = (dp[1]*k) - dp[0];
    }
    return dp[2]*k;
}

int main() {
    cout << numWays(3, 2) << "\n";
    return 0;
}