Редакция для Подотрезки с суммой, кратной k


Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
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. Алгоритм

  1. Считываем n и k.
  2. Создаём словарь cnt, где cnt[r] — сколько раз уже встречался остаток r у префиксных сумм.
  3. Изначально:
    • cnt[0] = 1
    • cur = 0
    • ans = 0
  4. Идём по массиву слева направо:
    • обновляем текущий остаток префиксной суммы cur;
    • добавляем к ответу cnt[cur], потому что каждый такой прошлый остаток даёт подходящий подотрезок;
    • увеличиваем cnt[cur] на 1.
  5. Выводим 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)

Комментарии

Еще нет ни одного комментария.