[l..r]
as a prefix [0..r] - [0..l]
.
For example, to find all subarrays [l..r]
whose sum of elements divides n
, we can think of this
as finding a subarray [l..r]
where the sum of elements modulo n
is zero.
This is CSES' subarray divisibiity problem:
int main() {
int n;
cin >> n;
vector<ll> xs(n);
for (int i = 0; i < n; ++i) {
cin >> xs[i]; xs[i] = xs[i] % n; if (xs[i] < 0) { xs[i] += n; }
}
ll count = 0; // number of subarrays with sum = 0 (mod n)
ll cursum = 0; // current sum [0..i]
// number of subarrays [0..r] (for some r) such that Σa[i] = count.
map<ll, ll> partial_sum_count;
partial_sum_count[0] = 1;
for (int i = 0; i < n; ++i) {
// current sum [0..i]
cursum = (cursum + xs[i]) % n;
// for each [0..j] (for j < i) with sum cursum, we want:
// sum([i..j]) = 0
// => sum([0..i]) - sum([0..j)) = 0
// => sum([0..i]) = sum([0..j))
// for each such `j`, we get one subarray.
auto it = partial_sum_count.find(cursum);
if (it != partial_sum_count.end()) {
count += it->second;
}
// partial sum [0..i] = cursum
partial_sum_count[cursum]++;
}
cout << count << "\n";
return 0;
}