Number of zigzag arrays i

Dynamic Programming Prefix Sum

32.9% 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)

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;
    // Shift values to 0..m-1 for indexing
    // dp[v][dir] = count of arrays of current length ending with value v and direction dir
    // dir=0: last step was up (prev < v)
    // dir=1: last step was down (prev > v)
    // For n=1: dp[v][0] = dp[v][1] = 1 (no direction yet, counted once each?)
    // Actually, for length 1, any valid starting value. No direction yet.
    // For length 2: a[0] != a[1]. Direction determined.
    // dp after position i:
    //   dp[v][0] = count of length-(i+1) arrays ending at value v with last step up
    //   dp[v][1] = count of length-(i+1) arrays ending at value v with last step down
    // Transition from position i to i+1:
    //   For new value w:
    //     If w > v: direction is up. Check: not 3 consecutive increasing → 
    //       last direction must NOT be up (otherwise a[i-1]<v<w = 3 increasing).
    //       So we can come from dp[v][1] (last step was down).
    //     If w < v: direction is down. Can come from dp[v][0] (last step was up).
    //     w == v: invalid.
    // 
    // dp_up[w] += sum over v < w: dp_down[v]
    // dp_down[w] += sum over v > w: dp_up[v]
    // Use prefix sums for efficiency.
    // Time: O(n * m).
    let mut dp_up = vec![0i64; m];   // ending at v, last step up
    let mut dp_down = vec![0i64; m]; // ending at v, last step down
    // Length 1: no direction yet, treat as both up and down possible (from virtual prev)
    // For length 2:
    // Actually, start with length 1 where all values have count 1 but no direction.
    // For transition to length 2:
    //   dp_up[w] = count where a[1]=w, a[0] < w = (w - offset) values of a[0]
    //   (since a[0] can be any value in [l,r] with no direction applied)
    //   But the direction constraint only applies from length 3+.
    // So init for length 2:
    //   dp_up[w] = w (shift: w values of a[0] < a[1]=w)
    //   dp_down[w] = m - 1 - w (values of a[0] > a[1]=w)
    // Then iterate from length 3 to n.
    for w in 0..m {
      dp_up[w] = w as i64 % MOD;        // w values less than w in [0..m-1]
      dp_down[w] = (m - 1 - w) as i64 % MOD; // m-1-w values greater than w
    }
    if n == 2 {
      return dp_up.iter().chain(dp_down.iter()).sum::<i64>() as i32 % MOD as i32;
    }
    for _ in 2..n {
      // Compute new_up[w] = sum over v < w: dp_down[v]  (prefix sum of dp_down)
      // Compute new_down[w] = sum over v > w: dp_up[v]   (suffix sum of dp_up)
      let mut prefix_down = vec![0i64; m + 1];
      for v in 0..m { prefix_down[v+1] = (prefix_down[v] + dp_down[v]) % MOD; }
      let mut suffix_up = vec![0i64; m + 1];
      for v in (0..m).rev() { suffix_up[v] = (suffix_up[v+1] + dp_up[v]) % MOD; }
      let mut new_up = vec![0i64; m];
      let mut new_down = vec![0i64; m];
      for w in 0..m {
        new_up[w] = prefix_down[w]; // sum dp_down[v] for v < w
        new_down[w] = suffix_up[w + 1]; // sum dp_up[v] for v > w
      }
      dp_up = new_up;
      dp_down = new_down;
    }
    let ans: i64 = dp_up.iter().chain(dp_down.iter()).sum::<i64>() % MOD;
    ans as i32
  }
}