Count non decreasing arrays with given digit sums

42.6% acceptance

Mar 31, 2026

You are given an integer array digitSum of length n. An array arr of length n is considered valid if: 0 <= arr[i] <= 5000 it is non-decreasing. the sum of the digits of arr[i] equals digitSum[i]. Return an integer denoting the number of distinct valid arrays. Since the answer may be large, return it modulo 109 + 7. An array is said to be non-decreasing if each element is greater than or equal to the previous element, if it exists.

Solution

Rust

Time O(n²)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn count_arrays(digit_sum: Vec<i32>) -> i32 {
    let m = 1_000_000_007i64;
    let n = digit_sum.len();
    
    // Precompute: for each digit sum s (0..50), list of values 0..5000 with that digit sum.
    // Actually, we need count of values in [lo, 5000] with digit sum s.
    // Better: precompute for each digit sum, the sorted list of valid values.
    
    let max_val = 5000;
    let max_ds = 50;
    
    // vals_by_ds[s] = sorted vec of values with digit sum s, 0..5000
    let mut vals_by_ds: Vec<Vec<i32>> = vec![Vec::new(); (max_ds + 1) as usize];
    for v in 0..=max_val {
      let ds = Self::dsum(v);
      if ds <= max_ds {
        vals_by_ds[ds as usize].push(v);
      }
    }
    
    // dp[v] = number of valid arrays ending with value v at current position.
    // But v can be up to 5000, so we track prefix sums.
    // dp_i[s] tracks: for digit_sum[i] = s, the values are vals_by_ds[s].
    // For non-decreasing: value at position i >= value at position i-1.
    // dp[i][v] = sum of dp[i-1][u] for all u <= v where dsum(u) = digit_sum[i-1].
    //
    // Approach: for each position, we have a set of valid values (those with correct digit sum).
    // We need cumulative counts.
    // Use an array of size 5001, cnt[v] = number of ways to have that value at current position.
    // Then for next position, for each valid value v', ways = sum of cnt[0..=v'].
    // Use prefix sum array.
    
    let s0 = digit_sum[0] as usize;
    if s0 > max_ds as usize || vals_by_ds[s0].is_empty() {
      return 0;
    }
    
    // cnt[v] for first position
    let mut cnt = vec![0i64; (max_val + 2) as usize];
    for &v in &vals_by_ds[s0] {
      cnt[v as usize] = 1;
    }
    
    for i in 1..n {
      let s = digit_sum[i] as usize;
      if s > max_ds as usize || vals_by_ds[s].is_empty() {
        return 0;
      }
      
      // Build prefix sum of cnt
      let mut prefix = vec![0i64; (max_val + 2) as usize];
      prefix[0] = cnt[0];
      for v in 1..=(max_val as usize) {
        prefix[v] = (prefix[v - 1] + cnt[v]) % m;
      }
      
      let mut new_cnt = vec![0i64; (max_val + 2) as usize];
      for &v in &vals_by_ds[s] {
        new_cnt[v as usize] = prefix[v as usize];
      }
      
      cnt = new_cnt;
    }
    
    let ans: i64 = cnt.iter().sum::<i64>() % m;
    ans as i32
  }
  
  fn dsum(mut v: i32) -> i32 {
    let mut s = 0;
    while v > 0 {
      s += v % 10;
      v /= 10;
    }
    s
  }
}