Subarray Sums Divisible by K

Problem

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

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

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
  • I can make it by using a prefix sum (after preparation, sum calculation become )
  • I realize that if sum of divisable by , and next to it is also divisable by , then is also divisable by , 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:
  • Space complexity:
  • Running time: 
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 , with counts the number of times remainder appears until this point.

  • At , your , you should look backward and see how many way to deduce from the to make it divisable
  • Thus, with added to the array and , it would makes more subarrays

Complexity

  • Time complexity:
  • Space complexity:
  • Running time:

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;
    }
};