[HackerRank][Python3] Construct the Array

Problem :

https://www.hackerrank.com/challenges/construct-the-array/problem

My Solution :

#!/usr/bin/env python3

def countArray(n, k, x):
    a, b = 1, 0
    m = 10**9 + 7
    for _ in range(2, n+1):
        a, b = (k-1)*b % m, ((k-2)*b + a) % m
    return a if x ==1 else b


n, k, x = map(int, input().strip().split())
answer = countArray(n, k, x)
print(answer)