Префиксные суммы + хешмап: подсчёт пар — задача «Good Subarrays» РАЗБОР

«Подсчитать количество подотрезков, удовлетворяющих условию» — формулировка, которая на полосе CF 1300–1800 встречается каждый второй контест. Когда условие алгебраическое (сумма, среднее, остаток), почти всегда работает один и тот же шаблон: записать условие через префиксные суммы, увидеть в нём равенство (или разность) двух значений, и свести задачу к подсчёту пар индексов с одинаковым (или отличающимся на константу) префиксом.

Эта задача — учебный экземпляр такого шаблона. Здесь условие «сумма подотрезка равна его длине» одной заменой переменной превращается в «сумма равна нулю», и дальше всё решается за один проход по массиву с хешмапом текущих частот префикса.

---

Что дано

В одном запуске несколько тестов. Для каждого теста:

• В первой строке — число $n$.
• Во второй — строка из $n$ символов; каждый символ — десятичная цифра $0$–$9$. Эта строка интерпретируется как массив $a_1, a_2, \ldots, a_n$.

Подотрезок $a_l, a_{l+1}, \ldots, a_r$ называется хорошим, если сумма его элементов равна длине подотрезка:

$$a_l + a_{l+1} + \cdots + a_r = r - l + 1.$$

Нужно вывести количество хороших подотрезков для каждого теста.

Ограничения

• $1 \le t \le 1000$ — число тестов.
• $1 \le n \le 10^5$ — длина массива в одном тесте.
• Сумма $n$ по всем тестам — до $10^5$.
• Цифры $a_i \in \{0, 1, \ldots, 9\}$.
• Время — 2 секунды, память — 256 МБ.

Формат ввода

В первой строке всех тестов — число $t$. Далее для каждого теста две строки: $n$ и строка из $n$ цифр.

Формат вывода

Для каждого теста — одно число, количество хороших подотрезков.

Пример 1

Вход:

3
3
120
5
11011
6
600005

Вывод:

3
6
1

Разберём первый тест руками

Массив $a = [1, 2, 0]$.

• Длина $1$: подотрезки $[1], [2], [0]$. Условие «сумма $=$ длина $= 1$»: $[1]$ — да; $[2]$ — нет; $[0]$ — нет. Хороших: $1$.
• Длина $2$: $[1, 2]$ сумма $3$ — нет; $[2, 0]$ сумма $2$ — да. Хороших: $1$.
• Длина $3$: $[1, 2, 0]$ сумма $3$ — да. Хороших: $1$.

Итого $3$ — сходится с ответом.

Решение «в лоб» — перебрать все $O(n^2)$ пар $(l, r)$, посчитать сумму на каждом и сравнить с длиной. При $n = 10^5$ это $10^{10}$ операций — на 4 порядка больше допустимого. Нужно $O(n)$ или $O(n \log n)$.

---

Идея решения

Перепишем условие в более удобной форме. Введём вспомогательный массив $b_i = a_i - 1$. Тогда

$$a_l + a_{l+1} + \cdots + a_r = (b_l + 1) + (b_{l+1} + 1) + \cdots + (b_r + 1) = (b_l + \cdots + b_r) + (r - l + 1).$$

Условие «сумма $a$ равна длине $r - l + 1$» эквивалентно

$$(b_l + b_{l+1} + \cdots + b_r) + (r - l + 1) = r - l + 1,$$

то есть

$$b_l + b_{l+1} + \cdots + b_r = 0.$$

Получили классическую постановку: подсчитать подотрезки массива $b$ с нулевой суммой.

Заведём префиксные суммы $P_0 = 0$, $P_i = b_1 + b_2 + \cdots + b_i$. Сумма $b$ на отрезке $[l, r]$ равна $P_r - P_{l-1}$. Условие «сумма равна нулю» эквивалентно $P_r = P_{l-1}$.

Значит, количество хороших подотрезков — это количество пар индексов $(i, j)$, $0 \le i < j \le n$, таких что $P_i = P_j$.

Подсчёт пар с одинаковым значением — стандартная задача на хешмап. Один проход:

1. Завести cnt: dict[int, int] = {0: 1} — стартовое значение $P_0 = 0$ уже встречено один раз.
2. Идти по массиву, обновлять текущий префикс $p$.
3. К ответу добавлять cnt[p] (столько раз раньше встречалось такое же значение префикса — каждая прошлая встреча образует пару с текущей).
4. Затем cnt[p] += 1.

После всего прохода в ответе накоплено число пар.

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

Биекция между «хорошими подотрезками $a$» и «парами равных префиксов $P_i = P_j$, $i < j$» — точная:

• $(l, r)$ хороший $\iff P_{l - 1} = P_r$ (по нашему алгебраическому выводу).
• Пара $(i, j)$, $i < j$, с $P_i = P_j$ задаёт подотрезок $(l = i + 1, r = j)$, который хороший.

Каждый хороший подотрезок задаёт ровно одну пару префиксов. Поэтому количество хороших подотрезков $=$ количество таких пар.

Подсчёт пар через «cnt до момента $j$» работает потому, что для каждого $j$ мы знаем, сколько индексов $i < j$ имели $P_i = P_j$ — это в точности cnt[P_j] до увеличения. После увеличения cnt[P_j] готов для следующего шага.

Размер значений префикса

$b_i = a_i - 1 \in \{-1, 0, 1, \ldots, 8\}$. Префикс может принимать значения от $-n$ до $8n$. При $n = 10^5$ — диапазон $[-10^5, 8 \cdot 10^5]$. Это влезает в int с большим запасом; хешмап с ключами в таком диапазоне — без проблем.

---

Решение: псевдокод

для каждого теста:
    прочитать n и строку s длины n
    cnt = {0: 1}                      // P_0 = 0 уже встречено
    p = 0                              // текущий префикс
    ans = 0
    для i от 1 до n:
        b_i = (s[i-1] как цифра) - 1
        p += b_i
        ans += cnt.get(p, 0)
        cnt[p] = cnt.get(p, 0) + 1
    вывести ans

Сложность: $O(n)$ в среднем на один тест.

---

Код решения

Различий между языками два значимых пункта. Первый — выбор контейнера: в Python dict или defaultdict(int) достаточно. В C++ — unordered_map<int, int>, но с защитой от анти-хеш-атак стоит использовать кастомный хеш или map<int, int>. Второй — чтение: строка цифр без разделителей читается одной строкой, преобразование в массив — поэлементно (ord(c) - ord('0') в Python, c - '0' в C++).

Комментарии по реализации.

• Тип ответа. $n = 10^5$, максимум пар — $\binom{n + 1}{2} \approx 5 \cdot 10^9$ при идеальной картине (все префиксы равны). int переполнится. В C++ — long long; в Python — встроенный int, без беспокойства.
• cnt[0] = 1 в начале. Аналог «нулевого префикса» — отрезки, начинающиеся с первого элемента, должны быть учтены. Без этого правильный ответ невозможен.
• Кастомный хеш в C++. На Codeforces встречаются тесты против стандартного unordered_map<int> (порождают $O(n)$ операций на вставку). SplitMix — стандартное решение: $x$ переписывается через серию умножений и сдвигов, распределение становится равномерным независимо от закономерностей входа.
• Чтение цифр из строки. Строка в условии — без разделителей; в C++ читается cin >> s как обычное слово; в Python — data[idx].decode() даёт str, по которому потом идём посимвольно.
• Глобальные структуры. В Python cnt создаётся в каждом тесте заново — это правильно, иначе тесты «загрязнят» друг друга. В C++ аналогично: unordered_map cnt; локально внутри while (t--).

---

Проверим на всех примерах из условия

Пример 1: 120

$a = [1, 2, 0]$, $b = [0, 1, -1]$. Префиксы $b$: $P = [0, 0, 1, 0]$ (длина $n + 1 = 4$).

Идём по $i = 1, 2, 3$:

Итого ans = 1 + 0 + 2 = 3. ✓

Пример 2: 11011

$a = [1, 1, 0, 1, 1]$, $b = [0, 0, -1, 0, 0]$. Префиксы: $P = [0, 0, 0, -1, -1, -1]$.

Итого: $1 + 2 + 0 + 1 + 2 = 6$. ✓

Пример 3: 600005

$a = [6, 0, 0, 0, 0, 5]$, $b = [5, -1, -1, -1, -1, 4]$. Префиксы: $P = [0, 5, 4, 3, 2, 1, 5]$.

Итого: $1$. ✓

Какой именно подотрезок? $P_1 = P_6 = 5$, значит хороший подотрезок — $a[2..6] = [0, 0, 0, 0, 5]$, длина $5$, сумма $5$. Подходит.

---

Крайние случаи, на которых можно споткнуться

1. Все элементы — единицы

$a = [1, 1, \ldots, 1]$, $b = [0, 0, \ldots, 0]$. Все префиксы равны нулю, и $P_0 = P_1 = \cdots = P_n = 0$. Хороших подотрезков $\binom{n + 1}{2} = \frac{n(n + 1)}{2}$ — каждый. Это согласуется с задачей: сумма любого подотрезка из одних единиц равна его длине.

При $n = 10^5$ ответ — около $5 \cdot 10^9$. long long обязателен.

2. Все элементы — нули

$a = [0, 0, \ldots, 0]$, $b = [-1, -1, \ldots, -1]$. Префиксы: $0, -1, -2, \ldots, -n$ — все разные. Хороших подотрезков нет, ответ $0$. Что согласуется: для длины $k$ нужна сумма $k$, но сумма нулей — $0$.

3. Длина $1$

Хороший подотрезок длины $1$ — это $a_i = 1$. Алгоритм видит: на шаге $i$ с $a_i = 1$ имеем $b_i = 0$, префикс не меняется, и предыдущее значение префикса было одним из ранее увиденных — встречаем минимум $P_{i-1}$, что даёт $+1$ к ответу.

4. Один тест с большим $n$

Сумма $n$ по тестам $\le 10^5$ — один тест с $n = 10^5$ ничем не отличается по сложности от $100$ тестов с $n = 1000$ каждый. Главное — не накладывать на каждый тест $O(n^2)$ инициализаций; локальный cnt в каждом тесте обнуляется автоматически.

5. Граница long long для ответа

Максимум — $\binom{10^5 + 1}{2} \approx 5 \cdot 10^9$. int (2^{31} \approx 2.1 \cdot 10^9) переполнится. long long обязателен в C++; в Python — нет проблем.

6. Хеш-атака на unordered_map

Без кастомного хеша возможен TL на специально подготовленных тестах. С SplitMix или аналогом — гарантированный $O(n)$ в среднем. Альтернатива — map<long long, long long>, $O(n \log n)$, медленнее, но устойчивее.

---

Типичные ошибки

1. Забытая инициализация cnt[0] = 1. Без этого подотрезки, начинающиеся с $l = 1$, не учитываются. Симптом: ответ меньше эталонного на конкретное число, зависящее от данных.
2. int вместо long long. Если массив длинный и все элементы — единицы, ответ выходит за пределы int. Лекарство — всегда long long под счётчик пар.
3. unordered_map<int, int> без защиты. Анти-хеш-тест даёт TL. Лекарство — кастомный хеш или map.
4. Подсчёт пар после увеличения cnt, а не до. Если сначала cnt[p] += 1, потом ans += cnt[p] — каждый префикс учтёт сам себя, в результате ans завышен на $n$.
5. Сброс cnt между тестами не происходит. Если cnt объявлен глобально и не очищается — соседние тесты «загрязняют» друг друга. Симптом: на одиночном тесте всё хорошо, в наборе — WA.
6. Чтение строки как массива чисел через пробел. В этой задаче строка цифр — без пробелов между ними. Чтение через cin >> n >> a[0] >> a[1] ... не сработает; нужно читать как строку и брать символы.
7. Перебор пар через двойной цикл. Соблазн «сначала собрать все префиксы, потом для каждой пары проверить» — $O(n^2)$ операций. Сразу подсчёт через хешмап — единственный правильный способ.

---

Анализ сложности

• Время: $O(n)$ на один тест (хешмап в среднем за $O(1)$ на операцию). Суммарно $O(\sum n) \le O(10^5)$.
• Память: $O(n)$ под хешмап в худшем случае (все префиксы разные).

---

Что ещё полезно потренировать

• Codeforces 1015E1 «Stars Drawing (Easy)» — другой вариант «префикс по строке + подсчёт», полоса CF $\sim 1500$.
• Codeforces 86D «Powerful Array» — алгоритм Мо с префиксами, шаг вперёд от текущей задачи; полоса CF $\sim 2200$, но идея «состояние через префиксы» здесь центральная.
• Codeforces 220B «Little Elephant and Array» — частоты в подотрезке через offline-обработку и префиксы; учит видеть, когда «префикс» — это не сумма, а другой накопительный признак.
• acmp.ru задача 442 «Подсчёт инверсий» — родственный шаблон «считать пары через структуру данных», вместо хешмапа здесь Fenwick.

---

Итого

• Идея: трансформация $b_i = a_i - 1$ превращает условие «сумма равна длине» в «сумма равна нулю»; равенство префиксов $P_i = P_j$ соответствует подотрезку с нулевой суммой; подсчёт пар через хешмап — $O(n)$.
• Сложность: $O(n)$ время и память на один тест.
• Ловушки: забытая cnt[0] = 1, int вместо long long, unordered_map без кастомного хеша, увеличение cnt до подсчёта вместо после.