Paint Fence

LintCode Q 514 - Paint Fence

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 color.
Return the total number of ways you can paint the fence.

Example 1:
Input: n=3, k=2 ; Output: 6
Explanation:
| | post1 | post2 | post3 |
|——|——-|——-|——-|
| way1 | 0 | 0 | 1 |
| way2 | 0 | 1 | 0 |
| way3 | 0 | 1 | 1 |
| way4 | 1 | 0 | 0 |
| way5 | 1 | 0 | 1 |
| way6 | 1 | 1 | 0 |
Example 2:
Input: n=2, k=2 ; Output: 4
Explanation:
| | post1 | post2 |
|——|——-|——-|
| way1 | 0 | 0 |
| way2 | 0 | 1 |
| way3 | 1 | 0 |
| way4 | 1 | 1 |
Notice: n and k are non-negative integers.

Solution : DP + Sliding Array

1. For post1: we have k choices;
2. For post2: we have k k* choices;
3. For post3:
   • if its color is same as post 2, then we have k (k - 1)* choices;
   • if its color is diff with post 2, then we have k k (k - 1) choices; Therefore the state transfer function is
     `dp[i] = (k - 1) * (dp[i - 1] + dp[i - 2]);`

At last we use the sliding array to optimze the space complexity.

Code:

public int numWays(int n, int k) {
	int[] dp = new int[]{0, k, k * k, 0};
	if (n == 1) return dp[1];
	if (n == 2) return dp[2];

	for (int i = 3; i <= n; i++) {
		dp[3] = (k - 1) * (dp[1] + dp[2]);
		dp[1] = dp[2];
		dp[2] = dp[3];
	}

	return dp[3];
}