Painting the Fence problem

The problem, simply

A fence has n posts in a row and you have k paint colors. Paint every post. The one rule: you may not have three or more posts next to each other in the same color. Two in a row is fine.

Give back how many different ways there are to paint the fence. The answer gets huge, so most versions ask for it modulo 1,000,000,007.

A story to remember it

Picture it

Examples

Way 1: the simple way

Look at the last post. Call T(i) the number of good paintings of i posts.

• If the last post has a new color, the first i - 1 posts can be any good painting, and you have k - 1 new colors. That gives (k - 1) x T(i - 1).
• If the last post copies the one before, then those two form a pair, and that pair must be a switch from post i - 2. That gives (k - 1) x T(i - 2).

So T(i) = (k - 1) x (T(i - 1) + T(i - 2)), with T(1) = k and T(2) = k x k. Write it as a plain recursive function.

class Solution {
    static final long MOD = 1_000_000_007L;

    public long countWays(int n, int k) {
        if (n == 0) return 0;
        if (n == 1) return k % MOD;
        if (n == 2) return (long) k * k % MOD;
        // new color after any painting of n-1, or a pair after any painting of n-2
        return (k - 1) * ((countWays(n - 1, k) + countWays(n - 2, k)) % MOD) % MOD;
    }
}

How long it takes: each call makes two calls, so the work grows like Fibonacci numbers, about O(2^n). Extra room: the call chain, O(n).

Way 2: the clever way

Walk the posts from left to right and keep two counts.

1. same: paintings whose last two posts match.
2. diff: paintings whose last post has a new color.
3. For the next post: newSame = diff (copying is only safe after a switch) and newDiff = (same + diff) x (k - 1) (switch after anything).
4. The total for the current post is same + diff.

class Solution {
    static final long MOD = 1_000_000_007L;

    public long countWays(int n, int k) {
        if (n == 0) return 0;
        long same = 0;          // last two posts match
        long diff = k;          // post 1: any color counts as a switch
        for (int i = 2; i <= n; i++) {
            long newSame = diff;                          // copy only after a switch
            long newDiff = (same + diff) % MOD * (k - 1) % MOD; // switch after anything
            same = newSame;
            diff = newDiff;
        }
        return (same + diff) % MOD;
    }
}

How long it takes: one short step per post, O(n). Extra room: two numbers, O(1).

Watch out for

• k = 1 should give 1 for one or two posts and 0 for three or more. The formulas handle it because k - 1 = 0.
• Take the modulo after every multiply. (same + diff) x (k - 1) can overflow a long without it.
• Read the rule carefully. Some versions (LeetCode 276) say "no more than two adjacent posts", which is this problem. A stricter version says no two adjacent posts may match at all, and then the answer is k x (k - 1)^(n - 1).

What interviewers ask next

• What if no three in a row means "at most m in a row"? Keep m counts: paintings ending with a run of length 1, 2, ..., m. Each new post either extends the run or starts a new one with k - 1 colors.
• Can you do it in O(log n)? Yes. The step from (same, diff) to the next pair is a fixed 2 by 2 multiply, so you can raise that 2 by 2 table to a power with repeated squaring.
• Paint houses where each color has a price and no two neighbors match. That is the "paint house" problem: keep the cheapest cost ending in each color.

Remember this

Grown-up words: the two counts are DP states, T(i) = (k - 1)(T(i-1) + T(i-2)) is a recurrence relation, keeping only the last values is space optimization, and the repeated squaring trick is matrix exponentiation.