Paint Fence Python Solution with Tests

lcpy gen -n 276   # by problem number
lcpy gen -s paint_fence   # by problem name

Input: n = 3, k = 2
Output: 6
Explanation: All the possibilities are shown.
Note that painting all the posts red or all the posts green is invalid because there cannot be three posts in a row with the same color.

Input: n = 1, k = 1
Output: 1

Input: n = 7, k = 2
Output: 42

class Solution:
    # Time: O(n)
    # Space: O(1)
    def num_ways(self, n: int, k: int) -> int:
        if n == 1:
            return k
        same, diff = k, k * (k - 1)
        for _ in range(n - 2):
            same, diff = diff, (same + diff) * (k - 1)
        return same + diff