Subarray Sum Divisible by K [Medium] — Prefix Sum Modulo

Sanjeev SharmaSanjeev Sharma
1 min read

Advertisement

Problem Statement

Return the number of non-empty subarrays with sum divisible by k.

Input: nums=[4,5,0,-2,-3,1], k=5Output: 7

Intuition

Prefix sum modulo k. If prefix[j] % k == prefix[i] % k, then subarray [i..j] sum is divisible by k. Count pairs using a HashMap of remainder frequencies.


Solutions

C++

int subarraysDivByK(vector<int>& nums, int k) {
    unordered_map<int,int> cnt; cnt[0]=1;
    int prefix=0, ans=0;
    for(int n:nums){
        prefix=(prefix+n%k+k)%k;
        ans+=cnt[prefix]++;
    }
    return ans;
}

Java

public int subarraysDivByK(int[] nums, int k) {
    Map<Integer,Integer> cnt=new HashMap<>(); cnt.put(0,1);
    int prefix=0, ans=0;
    for(int n:nums){
        prefix=((prefix+n)%k+k)%k;
        ans+=cnt.getOrDefault(prefix,0);
        cnt.merge(prefix,1,Integer::sum);
    }
    return ans;
}

Python

from collections import defaultdict

def subarraysDivByK(nums: list[int], k: int) -> int:
    cnt = defaultdict(int)
    cnt[0] = 1
    prefix = ans = 0
    for n in nums:
        prefix = (prefix + n) % k
        ans += cnt[prefix]
        cnt[prefix] += 1
    return ans

Complexity

  • Time: O(n), Space: O(k)

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro