Count distinct subarrays divisible by k in sorted array

Array Hash Table Prefix Sum

27.4% acceptance

Feb 24, 2026

You are given an integer array nums sorted in non-descending order and a positive integer k. A subarray of nums is good if the sum of its elements is divisible by k. Return an integer denoting the number of distinct good subarrays of nums. Subarrays are distinct if their sequences of values are.

Solution

Rust

Time O(n²)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn num_good_subarrays(nums: Vec<i32>, k: i32) -> i64 {
    // Key insight: since nums is sorted, subarray [l,r] is a duplicate of [l-1,r-1]
    // iff nums[l-1] == nums[r] (the element leaving equals the element entering).
    // So [l,r] is its "first occurrence" iff l==0 OR nums[l-1] != nums[r].
    //
    // For a fixed r, allowed l values are [0, lo_r] where lo_r = first index where
    // nums[lo_r] == nums[r]. This is because nums[l-1] == nums[r] for l > lo_r
    // (since nums is non-descending and all values at indices >= lo_r equal nums[r]).
    //
    // Algorithm (O(n)):
    // - Maintain a HashMap of prefix[l] % k for l in [0, lo_r].
    // - lo_r is non-decreasing, so we expand the map with a pointer.
    // - For each r, ans += map[prefix[r+1] % k].
    let n = nums.len();
    let k = k as i64;
    let mut prefix = vec![0i64; n + 1];
    for i in 0..n {
      prefix[i + 1] = prefix[i] + nums[i] as i64;
    }

    let mut map: std::collections::HashMap<i64, i64> = std::collections::HashMap::new();
    let mut ans = 0i64;
    let mut ptr = -1i64; // highest l index added to map so far
    let mut lo = 0usize; // first index of the current distinct value (lo_r)

    for r in 0..n {
      // Update lo_r: if value changed, lo jumps to r
      if r > 0 && nums[r] > nums[r - 1] {
        lo = r;
      }
      // Expand map: add prefix[l] % k for l in [ptr+1, lo]
      while ptr < lo as i64 {
        ptr += 1;
        *map.entry(prefix[ptr as usize] % k).or_insert(0) += 1;
      }
      // Count allowed l in [0, lo] with prefix[l] % k == prefix[r+1] % k
      let target = prefix[r + 1] % k;
      ans += map.get(&target).copied().unwrap_or(0);
    }
    ans
  }
}