Construct the array

construct_the_array.py

class Solution(object):

    def solve(self, n, k, x):

        a = [0] * n
        b = [0] * n

        if x == 1:
            a[0] = 1
        else:
            b[0] = 1

        for i in range(1, n):
            a[i] = b[i-1]
            b[i] = (a[i-1] * (k-1)) + (b[i-1] * (k-2))

        return a[n-1]


if __name__ == "__main__":
    n = 4
    k = 3
    x = 2

    ans = Solution().solve(n, k, x)
    print(ans)