Construct the array


Problem

Your goal is to find the number of ways to construct an array such that consecutive positions contain different values.

Specifically, we want to construct an array with nn elements such that each element between 11 and kk inclusive. We also want the first and last elements of the array to be 11 and xx

Example

For example, for n=4,k=3,x=2n=4, k=3, x=2 there are 33 ways, as shown here:

1 2 1 2
1 2 3 2
1 3 1 2

Constraints

  • 3n1053 \leq n \leq 10^5
  • 2k1052 \leq k \leq 10^5
  • 1xk1 \leq x \leq k

Submit your solution at here

Solution

Approach

We will apply the idea of Dynamic Programming.

  • Let f(i)f(i) indicates the number of ways to contruct array of ii length and the last number must not be xx.
  • Let g(i)g(i) indicates the number of ways to contruct array of ii length and the last number must be xx.

Easily we see the answer is g(n)g(n). We can calculate some initial value without much difficulty.

if x=1{f(1)=0f(2)=k1g(1)=1g(2)=0}\text{if } x = 1 \left \{ \begin{align*} f(1) &= 0\\ f(2) &= k - 1\\ g(1) &= 1\\ g(2) &= 0\\ \end{align*} \right \}
 if x1{f(1)=1f(2)=k2g(1)=0g(2)=1}\text{ if } x \neq 1 \left \{ \begin{align*} f(1) &= 1\\ f(2) &= k - 2\\ g(1) &= 0\\ g(2) &= 1\\ \end{align*} \right \}

From i=3i=3 onward, the formular stays consistent regardless of xx. f(i)f(i) and g(i)g(i) can be calculated using f(i1)f(i-1) and g(i1)g(i-1). There are only 1 way to pick xx at position ii so

g(i)=f(i1)g(i) = f(i - 1)

If the previous number is not xx we will lost 2 candidates. Thus, there are k2k - 2 ways to pick a number at ii If the previous number is xx we only lost 1 candidate. Thus, there are k1k - 1 ways to pick a number at ii Combine them altogether we have

f(i)=(k2)×f(i1)+(k1)×g(i1)f(i) = (k - 2) \times f(i - 1) + (k - 1) \times g(i - 1)

That’s conclude our solution.

Code

#define INF 1000000007
long countArray(int n, int k, int x) {
    long f[100001];
    long g[100001];
    if(x == 1) {
        g[1] = 1;
        f[1] = 0;
    } else {
        g[1] = 0;
        f[1] = 1;
    }
    for(int i=2;i<=n;i++) {
        f[i] = ((k-2)*f[i-1] + (k-1)*g[i-1])%INF;
        g[i] = f[i-1];
    }
    return g[n];
}

You can optimize it further by only store latest ff and gg, to get O(1)O(1) space complexity

#define INF 1000000007
long countArray(int n, int k, int x) {
    long f, g, tmp;
    if(x == 1) {
        g = 1;
        f = 0;
    } else {
        g = 0;
        f = 1;
    }
    for(int i=2;i<n;i++) {
        tmp = ((k-2)*f + (k-1)*g)%INF;
        g = f;
        f = tmp;
    }
    return f;
}