LeetCode 276: Paint Fence

Problem Restatement

We are given n fence posts and k colors.

Each post must be painted with exactly one color.

The rule is:

At most two adjacent fence posts may have the same color.

So this is allowed:

red, red, blue

But this is invalid:

red, red, red

We need to return the number of valid ways to paint all n posts. The source statement also notes that n and k are non-negative integers.

Input and Output

Function shape:

def numWays(n: int, k: int) -> int:
    ...

Examples

For:

n = 3
k = 2

Suppose the two colors are 0 and 1.

The valid ways are:

001
011
010
100
110
101

So the answer is:

6

For:

n = 2
k = 4

There are two posts.

The first post has 4 choices.

The second post also has 4 choices, because two adjacent posts may share the same color.

So the answer is:

4 * 4 = 16

First Thought: Brute Force

The brute force method tries every possible coloring.

There are k choices for each of the n posts, so the number of raw colorings is:

k ** n

For each coloring, we check whether it contains three consecutive posts with the same color.

That gives a correct solution, but it grows too quickly.

For example, if n = 20 and k = 5, there are:

5 ** 20

possible colorings to inspect.

We need to count valid arrangements without generating every arrangement.

Key Insight

The only thing that matters at the end of a partial painting is whether the last two posts have the same color.

So we keep two counts:

When we add a new post, there are two possibilities.

If the new post has the same color as the previous post, then the previous state must be diff.

Why?

Because if the previous state was already same, adding the same color again would create three equal adjacent posts.

So:

new_same = diff

If the new post has a different color from the previous post, then we can come from either same or diff.

There are k - 1 choices for the new color.

So:

new_diff = (same + diff) * (k - 1)

Algorithm

Handle small cases first.

If n == 0, there are no posts to paint:

return 0

If n == 1, there are k choices:

return k

For one post:

same = 0
diff = k

There cannot be two same adjacent posts yet, so same = 0.

There are k ways to paint the first post, so diff = k.

Then for every next post:

new_same = diff
new_diff = (same + diff) * (k - 1)

Update the two states.

At the end, the answer is:

same + diff

Correctness

After processing some number of posts, every valid painting belongs to exactly one of two groups.

The first group ends with two posts of the same color. These paintings are counted by same.

The second group ends with two posts of different colors. These paintings are counted by diff.

These two groups do not overlap, and together they cover all valid paintings.

When we paint the next post with the same color as the previous post, we may only extend a painting from diff. Extending a painting from same would create three adjacent posts with the same color, which violates the rule. Therefore new_same = diff.

When we paint the next post with a different color from the previous post, we may extend either kind of valid painting. There are k - 1 choices, because every color except the previous post’s color is allowed. Therefore new_diff = (same + diff) * (k - 1).

These transitions count every valid painting of the next length exactly once. Therefore, after all posts are processed, same + diff is the total number of valid ways.

Complexity

Implementation

class Solution:
    def numWays(self, n: int, k: int) -> int:
        if n == 0:
            return 0

        same = 0
        diff = k

        for _ in range(1, n):
            new_same = diff
            new_diff = (same + diff) * (k - 1)

            same = new_same
            diff = new_diff

        return same + diff

Code Explanation

We first handle the empty case:

if n == 0:
    return 0

Then we initialize the state for one post:

same = 0
diff = k

With one post, there are no adjacent posts yet, so same is 0.

There are k choices for the first post, so diff is k.

Then we add posts one by one:

for _ in range(1, n):

For each new post, if it has the same color as the previous post, the previous two posts must have been different:

new_same = diff

If it has a different color, we can extend any valid previous painting and choose one of k - 1 colors:

new_diff = (same + diff) * (k - 1)

Finally, we return both valid ending states:

return same + diff

Testing

def test_num_ways():
    s = Solution()

    assert s.numWays(0, 3) == 0
    assert s.numWays(1, 3) == 3
    assert s.numWays(2, 4) == 16
    assert s.numWays(3, 2) == 6
    assert s.numWays(3, 1) == 0
    assert s.numWays(4, 2) == 10

    print("all tests passed")

test_num_ways()

Test meaning: