Редакция для Сколько подотрезков с суммой, кратной k
Submitting an official solution before solving the problem yourself is a bannable offence.
1. Идея
Нужно посчитать количество подотрезков массива, сумма на которых делится на k.
Прямой перебор всех подотрезков слишком медленный: их O(n^2), а n может быть до 2 * 10^5.
Ключевая идея — перейти к префиксным суммам и смотреть не на сами суммы, а на их остатки по модулю k.
Если два префикса дают одинаковый остаток по модулю k, то сумма между ними делится на k.
2. Наблюдения
Обозначим pref[i] — сумму первых i элементов массива.
Тогда сумма подотрезка с позиции l по r равна:
pref[r + 1] - pref[l]
Эта сумма делится на k тогда и только тогда, когда
pref[r + 1] % k == pref[l] % k
Значит, задача сводится к такой:
- идём слева направо,
- считаем текущую префиксную сумму по модулю
k, - если такой остаток уже встречался раньше
tраз, то можно образоватьtподходящих подотрезков, заканчивающихся в текущей позиции.
Для хранения количества уже встреченных остатков удобно использовать словарь.
Ещё важное наблюдение: пустой префикс тоже нужно учесть. Его сумма равна 0, значит остаток 0 уже встречался один раз до начала массива.
Поэтому в начале делаем cnt[0] = 1.
3. Алгоритм
- Считать
n,kи массивa. - Создать словарь
cnt, где будет храниться, сколько раз встречался каждый остаток префиксной суммы. - Инициализировать:
cnt[0] = 1cur = 0— текущий остаток префиксной суммыans = 0— ответ
- Для каждого элемента
xмассива:- обновить остаток:
cur = (cur + x) % k - если такой остаток уже встречался, прибавить его количество к ответу
- увеличить
cnt[cur]на 1
- обновить остаток:
- Вывести
ans.
4. Почему это работает
Пусть в некоторый момент мы обработали элементы до позиции r.
cur — это остаток от суммы a[0] + a[1] + ... + a[r] по модулю k, то есть остаток префикса pref[r + 1].
Теперь рассмотрим любой более ранний префикс pref[l].
Если
pref[l] % k == pref[r + 1] % k
то разность
pref[r + 1] - pref[l]
делится на k.
А эта разность — как раз сумма подотрезка a[l] ... a[r].
Значит, каждый ранее встреченный такой же остаток даёт один подходящий подотрезок, заканчивающийся в r.
Поэтому, когда текущий остаток равен cur, нужно прибавить к ответу количество предыдущих префиксов с остатком cur.
После этого текущий префикс тоже становится доступным для следующих позиций, поэтому увеличиваем cnt[cur].
Так мы учитываем каждый подходящий подотрезок ровно один раз — в момент, когда обрабатываем его правую границу.
Почему нужно отдельно учитывать отрицательные числа
Элементы массива могут быть отрицательными. В C++ остаток от отрицательного числа тоже может быть отрицательным, поэтому после вычисления
cur = (cur + a[i]) % k
нужно сделать:
- если
cur < 0, тоcur += k
Это приводит остаток к стандартному диапазону от 0 до k - 1.
В Python оператор % уже даёт неотрицательный остаток при положительном k, поэтому дополнительная корректировка не нужна.
5. Сложность
Для каждого элемента выполняется константное число операций со словарём.
- Время:
O(n)в среднем - Память:
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;
vector<long long> a(n);
for (int i = 0; i < n; i++) {
cin >> a[i];
}
unordered_map<long long, long long> cnt;
cnt[0] = 1;
long long cur = 0;
long long ans = 0;
for (int i = 0; i < n; i++) {
cur = (cur + a[i]) % k;
if (cur < 0) {
cur += k;
}
if (cnt.find(cur) != cnt.end()) {
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)
Комментарии