Редакция для Сколько подотрезков с суммой, кратной 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.

Если два префикса дают одинаковый остаток по модулю 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. Алгоритм

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

Комментарии

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