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 elements such that each element between and inclusive. We also want the first and last elements of the array to be and

Example

For example, for there are ways, as shown here:

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

Constraints

Submit your solution at here

Solution

Approach

We will apply the idea of Dynamic Programming.

• Let indicates the number of ways to contruct array of length and the last number must not be .
• Let indicates the number of ways to contruct array of length and the last number must be .

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

From onward, the formular stays consistent regardless of . and can be calculated using and . There are only 1 way to pick at position so

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

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 and , to get 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;
}