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