Number of zigzag arrays ii

Math Dynamic Programming

57.5% acceptance

Feb 25, 2026

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

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

impl Solution {
  pub fn zig_zag_arrays(n: i32, l: i32, r: i32) -> i32 {
    const MOD: i64 = 1_000_000_007;
    let n = n as usize;
    let m = (r - l + 1) as usize; // <= 74
    // State vector: [dp_up[0..m], dp_down[0..m]] of length 2m.
    // Transition matrix T of size 2m x 2m:
    //   new_up[w] = sum_{v<w} dp_down[v]   → T[w][m+v] for v < w
    //   new_down[w] = sum_{v>w} dp_up[v]   → T[m+w][v] for v > w
    let sz = 2 * m;
    // Build transition matrix
    let mat_mul = |a: &Vec<Vec<i64>>, b: &Vec<Vec<i64>>| -> Vec<Vec<i64>> {
      let s = a.len();
      let mut c = vec![vec![0i64; s]; s];
      for i in 0..s {
        for k in 0..s {
          if a[i][k] == 0 { continue; }
          for j in 0..s {
            c[i][j] = (c[i][j] + a[i][k] * b[k][j]) % MOD;
          }
        }
      }
      c
    };
    let mat_pow = |mut base: Vec<Vec<i64>>, mut exp: usize| -> Vec<Vec<i64>> {
      let s = base.len();
      let mut result = vec![vec![0i64; s]; s];
      for i in 0..s { result[i][i] = 1; } // identity
      while exp > 0 {
        if exp & 1 == 1 { result = mat_mul(&result, &base); }
        base = mat_mul(&base, &base);
        exp >>= 1;
      }
      result
    };
    // Build transition matrix T
    let mut t = vec![vec![0i64; sz]; sz];
    // new_up[w] (index w) = sum_{v < w} dp_down[v] (index m+v)
    for w in 0..m {
      for v in 0..w {
        t[w][m + v] = 1;
      }
    }
    // new_down[w] (index m+w) = sum_{v > w} dp_up[v] (index v)
    for w in 0..m {
      for v in w+1..m {
        t[m + w][v] = 1;
      }
    }
    if n == 1 {
      return m as i32 % MOD as i32;
    }
    // Initial state for length 2: dp_up[w] = w, dp_down[w] = m-1-w
    let mut state = vec![0i64; sz];
    for w in 0..m {
      state[w] = w as i64; // dp_up[w]
      state[m + w] = (m - 1 - w) as i64; // dp_down[w]
    }
    if n == 2 {
      return (state.iter().sum::<i64>() % MOD) as i32;
    }
    // Apply transition matrix (n-2) times
    let t_pow = mat_pow(t, n - 2);
    // Multiply state by t_pow
    let mut new_state = vec![0i64; sz];
    for i in 0..sz {
      for j in 0..sz {
        new_state[i] = (new_state[i] + t_pow[i][j] * state[j]) % MOD;
      }
    }
    (new_state.iter().sum::<i64>() % MOD) as i32
  }
}