3700. Number of ZigZag Arrays II

3700. Number of ZigZag Arrays II

Problem Statement

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 109+710^9 + 7.

This problem is a continuation of 3699. Number of ZigZag Arrays I. The only difference is the constraints on n, which is now up to 10910^9, while the range r - l is at most 75.

Solution

This problem builds directly upon the dynamic programming solution for Number of ZigZag Arrays I. The core DP state and transitions remain the same, but the much larger value of n makes a linear-time DP approach too slow. We must find a way to optimize the calculation of the n-th DP state.

From DP to Matrix Form

Let’s recall the DP transitions. After normalizing the range [l, r] to [1, m] where m = r - l + 1, we have:

  • dp[i][j][0] (ending with ... < j): dp[i][j][0]=p=j+1mdp[i1][p][1]dp[i][j][0] = \sum_{p=j+1}^{m} dp[i-1][p][1]
  • dp[i][j][1] (ending with ... > j): dp[i][j][1]=p=1j1dp[i1][p][0]dp[i][j][1] = \sum_{p=1}^{j-1} dp[i-1][p][0]

This system of linear recurrences is a strong hint that we can use matrices. Let’s define a state vector V_i of size 2m that represents the DP values for a given length i: V_i = [dp[i][1][0], ..., dp[i][m][0], dp[i][1][1], ..., dp[i][m][1]]^T

Our goal is to find a 2m x 2m transition matrix M such that V_i = M * V_{i-1}. Looking at the dependencies, dp[i][...][0] only depends on dp[i-1][...][1], and dp[i][...][1] only depends on dp[i-1][...][0]. This structure suggests a block matrix for M:

M=(0AB0) M = \begin{pmatrix} 0 & A \\ B & 0 \end{pmatrix}

Here, 0 is an m x m zero matrix, and A and B are m x m matrices that handle the transitions.

  • Matrix A: This matrix transforms the “down” states (dp[i-1][...][1]) into the “up” states (dp[i][...][0]). The equation is dp[i][j][0] = sum_{p=j+1 to m} dp[i-1][p][1]. This means the j-th row of A should have 1s for all columns p where p > j. Ajk={1if k>j0otherwiseA_{jk} = \begin{cases} 1 & \text{if } k > j \\ 0 & \text{otherwise} \end{cases}

  • Matrix B: This matrix transforms the “up” states (dp[i-1][...][0]) into the “down” states (dp[i][...][1]). The equation is dp[i][j][1] = sum_{p=1 to j-1} dp[i-1][p][0]. This means the j-th row of B should have 1s for all columns p where p < j. Bjk={1if k<j0otherwiseB_{jk} = \begin{cases} 1 & \text{if } k < j \\ 0 & \text{otherwise} \end{cases}

For example, if m = 4, the matrices A and B would be: A=(0111001100010000),B=(0000100011001110) A = \begin{pmatrix} 0 & 1 & 1 & 1 \\ 0 & 0 & 1 & 1 \\ 0 & 0 & 0 & 1 \\ 0 & 0 & 0 & 0 \end{pmatrix}, \quad B = \begin{pmatrix} 0 & 0 & 0 & 0 \\ 1 & 0 & 0 & 0 \\ 1 & 1 & 0 & 0 \\ 1 & 1 & 1 & 0 \end{pmatrix}

Matrix Exponentiation

With the transition V_i = M * V_{i-1}, we can find V_n by starting from a base case and applying the matrix M repeatedly. V_n = M * V_{n-1} = M^2 * V_{n-2} = ... = M^{n-2} * V_2

Since n can be very large, we can’t perform n-2 matrix multiplications. Instead, we use Matrix Exponentiation (also known as binary exponentiation for matrices) to calculate Mn2M^{n-2} in O((2m)3logn)O((2m)^3 \log n) time. For a more detailed explanation of the algorithm, you can refer to this article on Matrix Exponentiation.

The base case is the state vector V_2, which represents arrays of length 2.

  • dp[2][j][0] (A[1] < A[2]=j): A[1] can be any value from l to l+j-2, so there are j-1 possibilities.
  • dp[2][j][1] (A[1] > A[2]=j): A[1] can be any value from l+j to r, so there are m-j possibilities.

So, the initial state vector V_2 is: V_2 = [m-1, m-2, ..., 0, 0, 1, ..., m-1]^T (Note: the provided code has a slight off-by-one in the DP base case, mapping [1,m] to dp[2][j][0] = m-j and dp[2][j][1] = j-1. The logic remains the same.)

After computing M=Mn2M' = M^{n-2}, we calculate the final state V_n = M' * V_2. The total number of ZigZag arrays is the sum of all elements in V_n, modulo 109+710^9 + 7.

View Code
static const auto fast_io = []() {
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(NULL);
    return 0;
}();
#define N 2005
class Solution {
public:
    int mod = 1e9 + 7;

    vector<vector<long long>> multi(vector<vector<long long>>& a,
                                    vector<vector<long long>>& b) {
        int n = a.size();
        vector<vector<long long>> c(n, vector<long long>(n, 0));
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                for (int k = 0; k < n; k++) {
                    c[i][j] = (c[i][j] + a[i][k] * b[k][j]) % mod;
                }
            }
        }
        return c;
    }

    vector<vector<long long>> power(vector<vector<long long>> a, int n) {
        int len = a.size();
        vector<vector<long long>> res(len, vector<long long>(len));
        for (int i = 0; i < len; i++)
            res[i][i] = 1;
        while (n) {
            if (n & 1)
                res = multi(res, a);
            a = multi(a, a);
            n /= 2;
        }
        return res;
    }

    int zigZagArrays(int n, int l, int r) {
        if (n == 1) {
            return r - l + 1;
        }
        if (n == 2) {
            long long m = r - l + 1;
            return (m * (m - 1)) % mod;
        }
        long long m = r - l + 1;
        vector<vector<long long>> tran(2 * m, vector<long long>(2 * m, 0));
        for (int i = 0; i < m; i++) {
            for (int j = i + 1; j < m; j++) {
                tran[i][j + m] = 1;
            }
        }
        for (int i = 1; i < m; i++) {
            for (int j = 0; j < i; j++) {
                tran[i + m][j] = 1;
            }
        }
        tran = power(tran, n - 2);
        vector<long long> start(2 * m, 0);
        for (int i = 0; i < m; i++) {
            start[i] = m - 1 - i;
            start[i + m] = i;
        }
        vector<long long> end(2 * m, 0);
        for (int i = 0; i < 2 * m; i++) {
            for (int j = 0; j < 2 * m; j++) {
                end[i] = (end[i] + tran[i][j] * start[j]) % mod;
            }
        }
        long long ans = 0;
        for (int i = 0; i < 2 * m; i++) {
            ans = (ans + end[i]) % mod;
        }
        return ans;
    }
};

Time and Space Complexity

  • Time Complexity: O((rl)3logn)O((r-l)^3 \log n). The size of the matrix is 2m x 2m where m = r-l+1. Matrix multiplication takes O(m3)O(m^3) and exponentiation takes O(logn)O(\log n) multiplications.
  • Space Complexity: O((rl)2)O((r-l)^2) to store the transition matrix.