CSES - Subarray Divisibility

Authors: Qi Wang, Brad Ma

Problem

We are asked to find the number of subarrays that are divisible by NN. In other words, we're supposed to find the number of subarrays with sum equal to 0(modN)0 \pmod N.

Explanation

Notice that any sum of a subarray can be represented as the difference of two prefixes.

First, let sum\texttt{sum} represent the prefix sum of array aa modulo NN.

With our prefix sums knowledge,

sum(i,j)=sum(0,j)sum(0,i1)\texttt{sum}(i, j) = \texttt{sum}(0, j) - \texttt{sum}(0, i-1)

Since we want to calculate the number of sum(i,j)\texttt{sum}(i, j) that equals to 0(modN)0\pmod N, sum(0,j)\texttt{sum}(0, j) must be equal to sum(0,i1)\texttt{sum}(0, i-1) for their difference to be 00.

Now, we calculate pmod[i]\texttt{pmod}[i], the number of prefixes with remainder equivalent to i(modN)i\pmod{N}. Then the number of pairs contributed by ii is

(pmod[i]2)=pmod[i](pmod[i]1)/2{\texttt{pmod}[i]\choose{2}} = \texttt{pmod}[i] \cdot (\texttt{pmod}[i] - 1) / 2

The answer is just the sum of this quantity over all ii.

Implementation

Time Complexity: O(N)\mathcal{O}(N)

C++

#include <iostream>
#include <vector>
using namespace std;
/**
* @author Qi Wang
* (detemplifying courtesy of Kevin Sheng)
*/
int main() {

Java

import java.io.*;
import java.util.*;
public class subarrayDivisibility {
public static void main(String[] args) {
Kattio io = new Kattio();
int n = io.nextInt();
long[] M = new long[n];

Python

n = int(input())
arr = map(int, input().split())
residue_counts = [0] * n
partial_sum = 0
residue_counts[partial_sum] = 1
for a in arr:
partial_sum += a
partial_sum = partial_sum % n
residue_counts[partial_sum] += 1
# each subarray with sum divisible by n corresponds to
# a pair of indices that have the same residue
print(sum(r * (r - 1) // 2 for r in residue_counts))

Join the USACO Forum!

Stuck on a problem, or don't understand a module? Join the USACO Forum and get help from other competitive programmers!