Количество подотрезков с заданной суммой — разбор

1. Идея

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

Перебирать все пары (l, r) напрямую нельзя: это заняло бы O(n^2), что слишком медленно при n до 2 * 10^5.

Главная идея — использовать префиксные суммы и хранить, сколько раз уже встречалась каждая префиксная сумма.

Пусть pref — сумма элементов от начала массива до текущей позиции. Тогда сумма подотрезка l..r равна:

prefix[r] - prefix[l - 1]

Чтобы эта сумма была равна S, должно выполняться:

prefix[l - 1] = prefix[r] - S

Значит, для каждого правого конца r надо узнать, сколько раз раньше встречалась префиксная сумма prefix[r] - S.

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

Наблюдение 1

Если известны префиксные суммы, то сумма любого подотрезка считается очень просто.

Например, если:

• prefix[i] — сумма первых i элементов,
• тогда сумма подотрезка l..r равна prefix[r] - prefix[l - 1].

Наблюдение 2

Для фиксированного конца r мы хотим найти количество таких l, что:

prefix[r] - prefix[l - 1] = S

Переносим: prefix[l - 1] = prefix[r] - S

То есть задача сводится к подсчёту: сколько раз раньше встречалось значение prefix[r] - S.

Наблюдение 3

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

Наблюдение 4

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

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

3. Алгоритм

Будем идти по массиву слева направо.

Поддерживаем:

• pref — текущую префиксную сумму,
• cnt[x] — сколько раз префиксная сумма x уже встретилась,
• ans — ответ.

Шаги:

1. Инициализируем cnt[0] = 1.
2. Для каждого элемента x:
   • прибавляем его к pref,
   • добавляем к ответу cnt[pref - S], если такое значение есть,
   • увеличиваем cnt[pref] на 1.
3. Выводим ans.

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

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

Пусть после обработки очередной позиции r текущая префиксная сумма равна pref.

Мы хотим узнать, сколько подотрезков заканчиваются в r и имеют сумму S.

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

a_l + ... + a_r = S

По определению префиксных сумм это эквивалентно:

prefix[r] - prefix[l - 1] = S

Значит:

prefix[l - 1] = prefix[r] - S

То есть каждому подходящему подотрезку l..r соответствует одно ранее встреченное значение префиксной суммы pref - S.

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

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

Таким образом:

• для каждого r учитываются все и только те подотрезки, которые заканчиваются в r и имеют сумму S;
• суммарно по всем r получаем точный ответ.

5. Сложность

Каждый элемент массива обрабатывается один раз, все операции со словарём или хеш-таблицей выполняются в среднем за O(1).

Итоговая сложность:

• время: 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[0] = 1;

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

    for (int i = 0; i < n; i++) {
        pref += a[i];
        if (cnt.find(pref - S) != cnt.end()) {
            ans += cnt[pref - S];
        }
        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)