Подотрезки с суммой S — разбор

1. Идея

Нужно посчитать количество непрерывных подотрезков массива, сумма которых равна S.

Если перебирать все пары границ подотрезка, получится слишком медленно: O(n^2), что не подходит при n до 2 * 10^5.

Главная идея — перейти от сумм подотрезков к префиксным суммам.

Пусть pref[i] — сумма первых i элементов массива. Тогда сумма подотрезка с границами l..r равна:

pref[r + 1] - pref[l]

Нам нужно, чтобы она была равна S, то есть:

pref[r + 1] - pref[l] = S

Отсюда:

pref[l] = pref[r + 1] - S

Значит, когда мы знаем текущую префиксную сумму, нужно понять, сколько раз раньше встречалась сумма pref - S. Это можно хранить в словаре.

2. Наблюдения

Наблюдение 1

Для каждого правого конца подотрезка достаточно знать, сколько подходящих левых концов существует.

Если текущая префиксная сумма равна pref, то количество подотрезков, заканчивающихся в текущей позиции и имеющих сумму S, равно числу уже встреченных префиксных сумм pref - S.

Наблюдение 2

Нужно заранее считать, что префиксная сумма 0 встречалась один раз.

Это соответствует пустому префиксу до начала массива. Благодаря этому корректно учитываются подотрезки, начинающиеся с первого элемента.

Например, если на некоторой позиции текущая префиксная сумма уже равна S, то подотрезок от начала массива до этой позиции должен быть засчитан. Именно для этого в словарь сразу кладём cnt[0] = 1.

Наблюдение 3

Элементы могут быть отрицательными, поэтому методы двух указателей или скользящего окна здесь не работают. Нужен именно подход с префиксными суммами и словарём.

3. Алгоритм

1. Считываем n, S и массив a.
2. Создаём словарь cnt, где cnt[x] — сколько раз уже встречалась префиксная сумма x.
3. Изначально записываем cnt[0] = 1.
4. Идём по массиву слева направо:
   • прибавляем текущий элемент к префиксной сумме pref;
   • добавляем к ответу cnt[pref - S], если такая сумма уже встречалась;
   • увеличиваем cnt[pref] на 1.
5. Выводим ответ.

4. Почему это работает

Докажем, что алгоритм считает ровно количество всех подотрезков с суммой S.

Рассмотрим текущую позицию r и текущую префиксную сумму pref = pref[r + 1].

Подотрезок l..r имеет сумму S тогда и только тогда, когда:

pref[r + 1] - pref[l] = S

Переносим:

pref[l] = pref[r + 1] - S

Значит, для фиксированного r каждый ранее встреченный префикс pref[l] = pref - S задаёт один подходящий подотрезок l..r.

Следовательно:

• число подходящих подотрезков, оканчивающихся в r, равно числу ранее встреченных префиксных сумм pref - S;
• именно это количество и добавляет алгоритм к ответу.

После этого мы заносим текущую префиксную сумму pref в словарь, чтобы она могла участвовать в подотрезках, которые заканчиваются позже.

Таким образом, каждый подходящий подотрезок:

• будет учтён ровно один раз — в момент обработки его правой границы;
• не будет пропущен.

Значит, алгоритм корректен.

5. Сложность

Для каждого элемента массива выполняется константное число операций со словарём.

• Время: O(n) в среднем
• Память: O(n)

6. Код на C++17

#include <iostream>
#include <vector>
#include <unordered_map>

using namespace std;

int main() {
    int n;
    long long S;
    cin >> n >> S;

    vector<long long> a(n);
    for (int i = 0; i < n; ++i) {
        cin >> a[i];
    }

    unordered_map<long long, long long> cnt;
    cnt.reserve((size_t)n * 2 + 10);
    cnt[0] = 1;

    long long pref = 0;
    long long ans = 0;

    for (int i = 0; i < n; ++i) {
        pref += a[i];
        auto it = cnt.find(pref - S);
        if (it != cnt.end()) {
            ans += it->second;
        }
        cnt[pref]++;
    }

    cout << ans << '\n';
    return 0;
}

7. Код на Python 3

n, S = map(int, input().split())
a = list(map(int, input().split()))

cnt = {0: 1}
pref = 0
ans = 0

for x in a:
    pref += x
    ans += cnt.get(pref - S, 0)
    cnt[pref] = cnt.get(pref, 0) + 1

print(ans)