3699. Number of Zigzag Arrays I

Problem Statement

Hard

You are given three integers n, l, and r.

A ZigZag array of length n is defined as follows:

Each element lies in the range [l, r].
No two adjacent elements are equal.
No three consecutive elements form a strictly increasing or strictly decreasing sequence.

Return the total number of valid ZigZag arrays.

Since the answer may be large, return it modulo 10⁹ + 7.

A sequence is said to be strictly increasing if each element is strictly greater than its previous one (if exists).

A sequence is said to be strictly decreasing if each element is strictly smaller than its previous one (if exists).

Example 1:

Input: n = 3, l = 4, r = 5

Output: 2

Explanation:

There are only 2 valid ZigZag arrays of length n = 3 using values in the range [4, 5]:

[4, 5, 4]
[5, 4, 5]

Example 2:

Input: n = 3, l = 1, r = 3

Output: 10

Explanation:

There are 10 valid ZigZag arrays of length n = 3 using values in the range [1, 3]:

[1, 2, 1], [1, 3, 1], [1, 3, 2]
[2, 1, 2], [2, 1, 3], [2, 3, 1], [2, 3, 2]
[3, 1, 2], [3, 1, 3], [3, 2, 3]

All arrays meet the ZigZag conditions.

Constraints:

3 <= n <= 2000
1 <= l < r <= 2000

Java

Source file

import java.util.Arrays;

class Solution {
    int MOD = 1_000_000_007, n, lt, rt, m;
    int[][][] memo;     // memo[dir][idx][val - lt]
    int[][][] prefix;   // prefix[dir][idx][t], t in [0, m]
    boolean[][] built;  // built[dir][idx]

    public int zigZagArrays(int n, int l, int r) {
        this.n = n;
        this.lt = l;
        this.rt = r;
        this.m = r - l + 1;

        memo = new int[2][n][m];
        for (int dir = 0; dir < 2; dir++)
            for (int idx = 0; idx < n; idx++)
                Arrays.fill(memo[dir][idx], -1);

        prefix = new int[2][n][m + 1];
        built = new boolean[2][n];

        int ans = 0;
        for (int val = lt; val <= rt; val++) {
            ans = (ans + solve(0, val, 1)) % MOD;
            ans = (ans + solve(0, val, 0)) % MOD;
        }
        return ans;
    }

    void ensurePrefix(int idx, int dir) {
        if (built[dir][idx]) return;
        int[] p = prefix[dir][idx];
        for (int val = lt; val <= rt; val++) {
            int v = solve(idx, val, dir);
            p[val - lt + 1] = (p[val - lt] + v) % MOD;
        }
        built[dir][idx] = true;
    }

    int solve(int idx, int val, int rising) {
        if (idx == n - 1)
            return 1;
        if (memo[rising][idx][val - lt] != -1)
            return memo[rising][idx][val - lt];

        int nextDir = rising ^ 1;
        ensurePrefix(idx + 1, nextDir);
        int[] p = prefix[nextDir][idx + 1];

        int ans;
        if (rising == 0) {
            ans = p[val - lt];
        } else {
            ans = (p[m] - p[val - lt + 1] + MOD) % MOD;
        }
        return memo[rising][idx][val - lt] = ans;
    }
}