Subarray Sums Divisible by K


Problem

Given an integer array numsnums and an integer kk, return the number of non-empty subarrays that have a sum divisible by kk.

A subarray is a contiguous part of an array.

Example

Input: nums = [4,5,0,-2,-3,1], k = 5
Output: 7
Explanation: There are 7 subarrays with a sum divisible by k = 5:
[4, 5, 0, -2, -3, 1], [5], [5, 0], [5, 0, -2, -3], [0], [0, -2, -3], [-2, -3]
Input: nums = [5], k = 9
Output: 0

Constraints

  • 1nums.length3×1041 \leq nums.length \leq 3 \times 10^4
  • 104numsi104-10^4 \leq nums_i \leq 10^4
  • 2k1042 \leq k \leq 10^4

Submit your solution at here

Solution

Naive brute force (AC, almost TLE)

  • First idea come to my mind is that I could traverse all subarray and count, the complexity is O(n3)O(n^3)
  • I can make it O(n2)O(n^2) by using a prefix sum (after O(n)O(n) preparation, sum calculation become O(1)O(1))
  • I realize that if sum of numsi,jnums_{i,j} divisable by kk, and numsj+1,xnums_{j+1,x} next to it is also divisable by kk, then numsi,xnums_{i,x} is also divisable by kk, I can cut off some calculation if the array happen to have consecutive segments like that

Complexity

Leetcode’s time limit is pretty generous so my naive solution luckily AC

  • Time complexity: O(n2)O(n^2)
  • Space complexity: O(n)O(n)
  • Running time: 2900ms2900ms
class Solution {
public:
    int subarraysDivByK(vector<int>& nums, int k) {
        int n = nums.size();
        vector<int> prefix(n+1, 0);
        for(int i = 1;i<=n;i++) {
            prefix[i] = prefix[i-1]+nums[i-1];
        }
        vector<bool> vis(n, false);
        int ret = 0;
        for(int i = 0;i<=n;i++) {
            if(vis[i]) {
                continue;
            }
            int segment = 0;
            int mid = i;
            for(int j = i+1;j<=n;j++) {
                int sum = prefix[j] - prefix[mid];
                if(sum%k == 0) {
                    segment++;
                    mid = j;
                    vis[j] = true;
                }
            }
            ret += (segment+1)*segment/2;
            vis[i] = true;
        }
        return ret;
    }
};

Intuition

Turned out the intended solution is so simple and elegant. You traverse the prefix sum array. Maintain array countcount, with countrcount_r counts the number of times remainder rr appears until this point.

  • At numsinums_i, your summodk=rsum\mod k = r, you should look backward and see how many way to deduce rr from the sumsum to make it divisable
  • Thus, with numsinums_i added to the array and summodk=rsum \mod k = r, it would makes countrcount_r more subarrays

Complexity

  • Time complexity: O(n)O(n)
  • Space complexity: O(n)O(n)
  • Running time: 49ms49ms

Code

class Solution {
public:
    int subarraysDivByK(vector<int>& nums, int k) {
        int n = nums.size();
        vector<int> prefix(n+1, 0);
        for(int i = 1;i<=n;i++) {
            prefix[i] = prefix[i-1]+nums[i-1];
        }
        int ret = 0;
        vector<int> remainderCount(k, 0);
        for(int i = 1;i<=n;i++) {
            int remainder = (k + (prefix[i] % k)) % k;// double mod to avoid negative remainder
            ret += remainderCount[remainder];
            remainderCount[remainder]++;
        }
        ret += remainderCount[0];
        return ret;
    }
};