Code Implementation

class Solution:
    def subarraysDivByK(self, nums, K):
        count = 0
        prefix_mod = 0
        mod_count = {0: 1}
        for num in nums:
            prefix_mod = (prefix_mod + num) % K
            # Handle negative mods
            prefix_mod = (prefix_mod + K) % K
            count += mod_count.get(prefix_mod, 0)
            mod_count[prefix_mod] = mod_count.get(prefix_mod, 0) + 1
        return count
      

class Solution {
public:
    int subarraysDivByK(vector<int>& nums, int K) {
        unordered_map<int, int> mod_count;
        mod_count[0] = 1;
        int prefix_mod = 0, count = 0;
        for (int num : nums) {
            prefix_mod = (prefix_mod + num) % K;
            if (prefix_mod < 0) prefix_mod += K;
            count += mod_count[prefix_mod];
            mod_count[prefix_mod]++;
        }
        return count;
    }
};
      

class Solution {
    public int subarraysDivByK(int[] nums, int K) {
        Map<Integer, Integer> modCount = new HashMap<>();
        modCount.put(0, 1);
        int prefixMod = 0, count = 0;
        for (int num : nums) {
            prefixMod = (prefixMod + num) % K;
            if (prefixMod < 0) prefixMod += K;
            count += modCount.getOrDefault(prefixMod, 0);
            modCount.put(prefixMod, modCount.getOrDefault(prefixMod, 0) + 1);
        }
        return count;
    }
}
      

var subarraysDivByK = function(nums, K) {
    let count = 0;
    let prefixMod = 0;
    let modCount = {0: 1};
    for (let num of nums) {
        prefixMod = (prefixMod + num) % K;
        if (prefixMod < 0) prefixMod += K;
        count += (modCount[prefixMod] || 0);
        modCount[prefixMod] = (modCount[prefixMod] || 0) + 1;
    }
    return count;
};
      

Problem Description

You are given an array of integers nums and an integer K. Your task is to find the total number of contiguous subarrays whose sum is divisible by K.

• Each subarray must be contiguous (elements are next to each other in nums).
• Elements can be positive, negative, or zero.
• Return the count of all such subarrays.

Constraints:

• 1 ≤ nums.length ≤ 3 × 10⁴
• -10⁴ ≤ nums[i] ≤ 10⁴
• 2 ≤ K ≤ 10⁴

Thought Process

The straightforward way to solve this problem is to consider every possible subarray, compute its sum, and check if it is divisible by K. However, this brute-force approach would be too slow for large arrays.

To optimize, we need to avoid recalculating sums for overlapping subarrays. A key insight is to use prefix sums: the sum of elements from the start up to a given index. If two prefix sums have the same remainder when divided by K, the subarray between them is divisible by K. This is because the difference between the two prefix sums is a multiple of K.

By tracking how often each remainder occurs, we can efficiently count valid subarrays. This shifts our approach from checking every subarray to counting the frequency of prefix sum remainders.

Solution Approach

• Prefix Sum and Modulo: Compute a running sum (prefix sum) as we iterate through the array. At each step, calculate the remainder of the prefix sum divided by K.
• Hash Map for Remainders: Maintain a hash map (or array) that counts how many times each remainder has occurred so far.
• Counting Subarrays: If the current remainder has been seen before, it means there are as many subarrays ending at the current index whose sum is divisible by K as the count of that remainder so far. Add this count to the answer.
• Adjustment for Negative Numbers: Since modulo can be negative in some languages, always adjust the remainder to be in the range [0, K-1].
• Initialization: Initialize the hash map with 0:1 to handle subarrays starting from the beginning.

This approach ensures each element is processed only once, and all lookups and updates are constant time.

Example Walkthrough

Input: nums = [4,5,0,-2,-3,1], K = 5

• Initialize prefix_mod = 0, mod_count = {0:1}, count = 0
• Index 0: 4 → prefix_mod = (0 + 4) % 5 = 4
mod_count[4] not found, count stays 0. Increment mod_count[4] to 1.
• Index 1: 5 → prefix_mod = (4 + 5) % 5 = 4
mod_count[4] = 1, so add 1 to count (now 1). Increment mod_count[4] to 2.
• Index 2: 0 → prefix_mod = (4 + 0) % 5 = 4
mod_count[4] = 2, add 2 to count (now 3). Increment mod_count[4] to 3.
• Index 3: -2 → prefix_mod = (4 + -2) % 5 = 2
mod_count[2] not found, count stays 3. Increment mod_count[2] to 1.
• Index 4: -3 → prefix_mod = (2 + -3) % 5 = -1 % 5 = 4
  Adjust to positive: 4. mod_count[4] = 3, add 3 to count (now 6). Increment mod_count[4] to 4.
• Index 5: 1 → prefix_mod = (4 + 1) % 5 = 0
mod_count[0] = 1, add 1 to count (now 7). Increment mod_count[0] to 2.

Final answer: 7

Time and Space Complexity

• Brute-force Approach:
  For every subarray, sum its elements and check divisibility. There are O(N²) subarrays, and each sum takes O(1) to O(N) time. So time complexity is O(N²), which is too slow for large N.
• Optimized Approach:
  We process each element once, and hash map operations are O(1) on average.
  Time Complexity: O(N), where N is the length of nums.
  Space Complexity: O(K), for storing up to K different remainders in the hash map.

Summary

By leveraging prefix sums and a hash map to track remainders, we efficiently count all subarrays whose sum is divisible by K. The key insight is recognizing that matching remainders indicate subarrays with sums divisible by K. This approach reduces the problem from quadratic to linear time, making it both elegant and practical for large inputs.