Paint Fence Solution

paint_fence.py

# 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.

from functools import lru_cache

class Solution:
    @lru_cache(maxsize=None)
    def ways(self, a, b, p, n, k):
        if p == n:
            return 1
        ctr = 0
        for i in range(k):
            if a == -1 or b == -1 or len({a, b, i}) > 1:
                ctr += self.ways(b, i, p + 1, n, k)
        return ctr

    def num_ways(self, n: int, k: int) -> int:
        return self.ways(-1, -1, 0, n, k)