Given three positive integers n, k and x. The task is to count the number of different array that can be formed of size n such that each element is between 1 to k and two consecutive element are different. Also, the first and last elements of each array should be 1 and x respectively.
Input : n = 4, k = 3, x = 2 Output : 3
The idea is to use Dynamic Programming and combinatorics to solve the problem.
First of all, notice that the answer is same for all x from 2 to k. It can easily be proved. This will be useful later on.
Let the state f(i) denote the number of ways to fill the range [1, i] of array A such that A1 = 1 and Ai ≠ 1.
Therefore, if x ≠ 1, the answer to the problem is f(n)/(k – 1), because f(n) is the number of way where An is filled with a number from 2 to k, and the answer are equal for all such values An, so the answer for an individual value is f(n)/(k – 1).
Otherwise, if x = 1, the answer is f(n – 1), because An – 1 ≠ 1, and the only number we can fill An with is x = 1.
Now, the main problem is how to calculate f(i). Consider all numbers that Ai – 1 can be. We know that it must lie in [1, k].
- If Ai – 1 ≠ 1, then there are (k – 2)f(i – 1) ways to fill in the rest of the array, because Ai cannot be 1 or Ai – 1 (so we multiply with (k – 2)), and for the range [1, i – 1], there are, recursively, f(i – 1) ways.
- If Ai – 1 = 1, then there are (k – 1)f(i – 2) ways to fill in the rest of the array, because Ai – 1 = 1 means Ai – 2 ≠ 1 which means there are f(i – 2)ways to fill in the range [1, i – 2] and the only value that Ai cannot be 1, so we have (k – 1) choices for Ai.
By combining the above, we get
f(i) = (k - 1) * f(i - 2) + (k - 2) * f(i - 1)
This will help us to use dynamic programming using f(i).
Below is the implementation of this approach: