Choose numbers from two arrays in range

use std::collections::HashMap;

impl Solution {
  pub fn count_subranges(nums1: Vec<i32>, nums2: Vec<i32>) -> i32 {
    const MOD: i64 = 1_000_000_007;
    let n = nums1.len();
    let mut ans: i64 = 0;

    // For each starting index, extend to the right tracking diff = sum1 - sum2.
    // At each position i, pick nums1[i] (add to sum1) or nums2[i] (add to sum2).
    // We want diff == 0.
    // Use DP with HashMap on diff.
    // dp[diff] = number of ways to reach this diff for subarrays ending at current position.
    
    let mut dp: HashMap<i32, i64> = HashMap::new();
    
    for i in 0..n {
      let mut new_dp: HashMap<i32, i64> = HashMap::new();
      // Start a new subarray at i
      *new_dp.entry(nums1[i]).or_insert(0) += 1;
      *new_dp.entry(-nums2[i]).or_insert(0) += 1;
      
      // Extend existing subarrays
      for (&diff, &cnt) in &dp {
        // Pick nums1[i]
        *new_dp.entry(diff + nums1[i]).or_insert(0) = 
          (new_dp.get(&(diff + nums1[i])).unwrap_or(&0) + cnt) % MOD;
        // Pick nums2[i]
        *new_dp.entry(diff - nums2[i]).or_insert(0) = 
          (new_dp.get(&(diff - nums2[i])).unwrap_or(&0) + cnt) % MOD;
      }
      
      ans = (ans + new_dp.get(&0).unwrap_or(&0)) % MOD;
      dp = new_dp;
    }

    ans as i32
  }
}