Редакция для Подотрезки с суммой, кратной k
Submitting an official solution before solving the problem yourself is a bannable offence.
1. Идея
Нужно посчитать количество подотрезков массива, сумма на которых делится на k.
Прямой перебор всех подотрезков слишком медленный: их O(n^2), а n может быть до 2 * 10^5.
Главная идея — перейти от сумм подотрезков к префиксным суммам по модулю k.
Если обозначить:
pref[i]— сумма первыхiэлементов,- тогда сумма подотрезка
l..rравнаpref[r] - pref[l - 1].
Эта сумма делится на k тогда и только тогда, когда остатки pref[r] % k и pref[l - 1] % k равны.
Значит, задача сводится к подсчёту количества пар одинаковых остатков префиксных сумм.
2. Наблюдения
Наблюдение 1
Подотрезок l..r имеет сумму, кратную k, если:
a_l + a_{l+1} + ... + a_r делится на k.
Через префиксные суммы это:
pref[r] - pref[l - 1] делится на k.
А это эквивалентно тому, что:
pref[r] % k == pref[l - 1] % k.
Наблюдение 2
Если при просмотре массива слева направо текущий остаток префиксной суммы равен cur, то каждый предыдущий префикс с таким же остатком образует один подходящий подотрезок.
То есть если остаток cur уже встречался t раз, то текущий префикс добавляет в ответ t.
Наблюдение 3
Нужно не забыть про пустой префикс.
До начала массива сумма равна 0, и её остаток тоже 0. Это важно, потому что подотрезки, начинающиеся с первого элемента, тоже должны учитываться.
Поэтому изначально записываем:
- остаток
0встречался1раз.
Наблюдение 4
Элементы массива могут быть отрицательными.
В Python выражение (cur + x) % k уже даёт неотрицательный остаток.
В C++ остаток от отрицательного числа может быть отрицательным, поэтому используется нормализация:
((cur + x) % k + k) % k
3. Алгоритм
- Считываем
nиk. - Создаём словарь
cnt, гдеcnt[r]— сколько раз уже встречался остатокrу префиксных сумм. - Изначально:
cnt[0] = 1cur = 0ans = 0
- Идём по массиву слева направо:
- обновляем текущий остаток префиксной суммы
cur; - добавляем к ответу
cnt[cur], потому что каждый такой прошлый остаток даёт подходящий подотрезок; - увеличиваем
cnt[cur]на1.
- обновляем текущий остаток префиксной суммы
- Выводим
ans.
4. Почему это работает
Докажем корректность.
Пусть pref[i] — сумма первых i элементов, причём pref[0] = 0.
Рассмотрим произвольный подотрезок l..r. Его сумма равна:
pref[r] - pref[l - 1].
Этот подотрезок подходит тогда и только тогда, когда его сумма делится на k, то есть:
pref[r] - pref[l - 1] кратно k.
Это равносильно тому, что остатки pref[r] и pref[l - 1] по модулю k совпадают.
Значит, каждому подходящему подотрезку соответствует пара индексов префиксов (l - 1, r) с одинаковыми остатками. И наоборот, любая такая пара задаёт подходящий подотрезок.
Теперь посмотрим, что делает алгоритм.
Когда мы дошли до очередного префикса с остатком cur, словарь cnt уже хранит количество всех предыдущих префиксов с каждым возможным остатком. Если cnt[cur] = t, то существует ровно t значений начала подотрезка, для которых текущий конец образует сумму, кратную k.
Поэтому мы прибавляем t к ответу.
После этого увеличиваем cnt[cur], чтобы текущий префикс мог участвовать в подотрезках, заканчивающихся позже.
Так мы учитываем все подходящие подотрезки ровно один раз.
5. Сложность
- Время:
O(n)в среднем, потому что каждый элемент обрабатывается один раз, а операции со словарём выполняются заO(1)в среднем. - Память:
O(n)в худшем случае, если все остатки префиксных сумм разные.
6. Код на C++17
#include <iostream>
#include <vector>
#include <unordered_map>
using namespace std;
int main() {
int n;
long long k;
cin >> n >> k;
unordered_map<long long, long long> cnt;
cnt.reserve((size_t)n * 2 + 10);
long long cur = 0;
long long ans = 0;
cnt[0] = 1;
for (int i = 0; i < n; i++) {
long long x;
cin >> x;
cur = ((cur + x) % k + k) % k;
ans += cnt[cur];
cnt[cur]++;
}
cout << ans << '\n';
return 0;
}
7. Код на Python 3
n, k = map(int, input().split())
a = list(map(int, input().split()))
cnt = {0: 1}
cur = 0
ans = 0
for x in a:
cur = (cur + x) % k
ans += cnt.get(cur, 0)
cnt[cur] = cnt.get(cur, 0) + 1
print(ans)
Комментарии