CSES Solutions - Subarray Divisibility

Given an array arr[] of N integers, your task is to count the number of subarrays where the sum of values is divisible by N.

Examples:

Input: N = 5, arr[] = {3, 1, 2, 7, 4}
Output: 1
Explanation: There is only 1 subarray with sum divisible by 5, subarray {1, 2, 7}, sum = 10 and 10 is divisible by 5.

Input: N = 5, arr[] = {1, 2, 3, 4, 5}
Output: 4
Explanation: There are 4 subarrays with sum divisible by 5

• Subarray {5}, sum = 5 and 5 is divisible by 5.
• Subarray {2, 3}, sum = 5 and 5 is divisible by 5.
• Subarray {1, 2, 3, 4}, sum = 10 and 10 is divisible by 5.
• Subarray {1, 2, 3, 4, 5}, sum = 15 and 15 is divisible by 5.

Approach: To solve the problem, follow the below idea:

The problem is similar to Subarray Sums II. We can maintain a Map, that stores the (prefix sums % N) along with the number of times they have occurred. At any index i, let's say (prefix sums % N) = R, that is (sum of subarray arr[0...i] % N) = R. Now if there is an index j (j < i) such that (sum of subarray arr[0...j] % N) = R, then the remainder of subarray arr[j+1...i] will be equal to 0, which means that the subarray arr[j+1...i] is divisible by N. Therefore, for every index i, if we can find the count of prefixes before i which have remainder as arr[0...i], we can add the count to our answer and the sum of count for all indices will be the final answer.

This allows us to find a subarray within our current array (by removing a prefix from our current prefix) that leaves remainder = 0, when divided by N. Also, as we iterate through the array, we continuously update the map with the new remainders after each step so that all possible remainders are counted in the map as we traverse the array.

Step-by-step algorithm:

• Maintain a map, say remaindersCnt to store the count of occurrences of each (prefix sum % N).
• Maintain a variable remainder = 0, to calculate the remainder of prefix sum till any index and a variable cnt = 0 to count the number of subarrays with prefix sum divisible by N.
• Initialize remaindersCnt[0] = 1 so when we get any subarray with sum divisible by N, we can add remaindersCnt[pref % N] = remaindersCnt[0] = 1 to the answer.
• Iterate over all the elements arr[i],
  • Store the remainder of prefix sums of subarray arr[0...i] into remainder.
  • Add the frequency of remaindersCnt[remainder] to the cnt.
  • Increment the frequency of remainder by 1.
• Return the final answer as cnt.

Below is the implementation of the above algorithm:

#include <bits/stdc++.h>
#define ll long long int
using namespace std;

// Function to count the number of subarrays divisible by N
ll solve(vector<ll>& arr, ll N)
{
    // Map to store the frequency of prefix sums % N
    map<ll, ll> remaindersCnt;

    remaindersCnt[0] += 1;
    ll remainder = 0;
    ll cnt = 0;

    // Iterate over all the index and add the count of
    // subarrays with sum divisible by N
    for (int i = 0; i < N; ++i) {
        // Since arr[i] can be negative, we add N to the
        // remainder to avoid negative remainders
        remainder = ((remainder + arr[i]) % N + N) % N;
        cnt += remaindersCnt[remainder];
        remaindersCnt[remainder] += 1;
    }
    return cnt;
}

int main()
{
    // Sample Input
    ll N = 5;
    vector<ll> arr = { 1, 2, 3, 4, 5 };
  
    cout << solve(arr, N);
    return 0;
}

import java.util.HashMap;
import java.util.Map;

public class Main {
    public static int GFG(int[] arr, int N) {
        // Map to store the frequency of the prefix sums % N
        Map<Integer, Integer> remaindersCnt = new HashMap<>();
        remaindersCnt.put(0, 1);
        int remainder = 0;
        int cnt = 0;
        for (int i = 0; i < N; ++i) {
            // Since arr[i] can be negative and we add N to remainder to avoid negative remainders
            remainder = ((remainder + arr[i]) % N + N) % N;
            cnt += remaindersCnt.getOrDefault(remainder, 0);
            remaindersCnt.put(remainder, remaindersCnt.getOrDefault(remainder, 0) + 1);
        }
        return cnt;
    }

    public static void main(String[] args) {
        // Sample Input
        int N = 5;
        int[] arr = {1, 2, 3, 4, 5};
        System.out.println(GFG(arr, N));
    }
}

using System;
using System.Collections.Generic;

public class Solution
{
    public static int GFG(int[] arr, int N)
    {
        // Dictionary to store the frequency of the prefix sums % N
        Dictionary<int, int> remaindersCnt = new Dictionary<int, int>();
        remaindersCnt.Add(0, 1);
        int remainder = 0;
        int cnt = 0;
        for (int i = 0; i < N; ++i)
        {
            // Since arr[i] can be negative and we add N to remainder to avoid negative remainders
            remainder = ((remainder + arr[i]) % N + N) % N;
            cnt += remaindersCnt.ContainsKey(remainder) ? remaindersCnt[remainder] : 0;
            if (remaindersCnt.ContainsKey(remainder))
                remaindersCnt[remainder] = remaindersCnt[remainder] + 1;
            else
                remaindersCnt.Add(remainder, 1);
        }
        return cnt;
    }

    public static void Main(string[] args)
    {
        // Sample Input
        int N = 5;
        int[] arr = { 1, 2, 3, 4, 5 };
        Console.WriteLine(GFG(arr, N));
    }
}

function GFG(arr, N) {
    // Map to store the frequency of the prefix sums % N
    let remaindersCnt = new Map();
    remaindersCnt.set(0, 1);
    let remainder = 0;
    let cnt = 0;
    for (let i = 0; i < N; ++i) {
        // Since arr[i] can be negative and we add N to 
        // remainder to avoid negative remainders
        remainder = ((remainder + arr[i]) % N + N) % N;
        cnt += (remaindersCnt.get(remainder) || 0);
        remaindersCnt.set(remainder, (remaindersCnt.get(remainder) || 0) + 1);
    }
    return cnt;
}
// Sample Input
const N = 5;
const arr = [1, 2, 3, 4, 5];
console.log(GFG(arr, N));

# Function to count the number of subarrays divisible by N
def solve(arr, N):
    # Dictionary to store the frequency of prefix sums % N
    remainders_cnt = {0: 1}
    remainder = 0
    cnt = 0

    # Iterate over all the indices and add the count of
    # subarrays with sum divisible by N
    for i in range(len(arr)):
        # Since arr[i] can be negative, we add N to the
        # remainder to avoid negative remainders
        remainder = ((remainder + arr[i]) % N + N) % N
        cnt += remainders_cnt.get(remainder, 0)
        remainders_cnt[remainder] = remainders_cnt.get(remainder, 0) + 1
    
    return cnt

if __name__ == "__main__":
    # Sample Input
    N = 5
    arr = [1, 2, 3, 4, 5]

    print(solve(arr, N))

4

Time Complexity: O(N * logN), where N is the size of arr[].
Auxiliary Space: O(N)