Code:
public void solve(long[] arr, int n) {
long[] prefixSum = new long[n + 1];
prefixSum[0] = 0;
prefixSum[1] = arr[0] % n;
for (int i = 1; i < n; i++) {
prefixSum[i + 1] = (arr[i] + prefixSum[i]);
}
for (int i = 0; i <= n; i++) {
prefixSum[i] = prefixSum[i] % n;
}
long[] buckets = new long[n];
for (int i = 0; i <= n; i++) {
System.out.print(prefixSum[i] + " ");
buckets[(int)Math.abs(prefixSum[i])]++;
}
//System.out.println();
long res = 0;
for (int i = 0; i < n; i++) {
//System.out.print(buckets[i] + " ");
if (buckets[i] >= 2)
res += calcCombination(buckets[i], 2);
}
//System.out.println();
System.out.println(res);
}
// calc: nCk
private long calcCombination(long n, long k) {
if (k > n - k) {
k = n - k;
}
long res = 1;
for (int i = 0; i < k; i++) {
res *= (n - i);
res /= (i + 1);
}
return res;
}