LeetCode 276. Paint Fence

#include <iostream>
#include <vector>
using namespace std;

/*
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 colors.
Return the total number of ways you can paint the fence.
Note: n and k are non-negative integers.
*/

/*
It is better to use an example to illustrate the question.
if n == 0, return 0 method.
if n == 1, return k method (1 post can have k colors)

// k-1 is the colors for second post different from the first one.
// 1 is the color that  same as first post.
if n == 2, return k * (k - 1) + k * 1.

// (k * 1 * (k-1) + k * (k - 1) * (k - 1)) is the choice that different from second post.
// (k * (k-1) + k * 1) * 1, is the choice that same color as second post.
if n == 3, return (k * 1 * (k-1) + k * (k - 1) * (k - 1)) + (k * (k-1) + k * 1) * 1
*/
int numOfWays(int n, int k) {
if(n == 0) return 0;
if(n == 1) return k;
int same_color = k;
int different_color = k * (k - 1);
for(int i = 3; i <= n; ++i) {
int tmp = different_color;
different_color = different_color * (k-1) + same_color * (k-1);
same_color = tmp * 1;
}
return different_color + same_color;
}

int main(void) {
cout << numOfWays(3, 2) << endl;
}