560. Subarray Sum Equals K — разбор задачи

1. Условие задачи

Дано: массив целых чисел nums и целое число k.

Нужно: посчитать, сколько в массиве непрерывных непустых подмассивов с суммой элементов, равной k.

Подмассив обязан состоять из соседних элементов. Выбирать элементы через один нельзя.

Например:

nums = [1, 1, 1]
k = 2

Подходящих подмассивов два:

[1, 1, 1]
 └──┘       индексы 0..1
    └──┘    индексы 1..2

Поэтому ответ равен 2.

Ещё один пример:

nums = [1, 2, 3]
k = 3

Здесь подходят [1, 2] и [3], поэтому ответ снова равен 2.

Ограничения:

длина nums — от 1 до 2 * 10^4;
каждый элемент лежит в диапазоне от -1000 до 1000;
k лежит в диапазоне от -10^7 до 10^7;
в массиве могут быть нули и отрицательные числа.

Последний пункт особенно важен: обычное скользящее окно здесь не работает. Если справа добавить отрицательное число, сумма окна может уменьшиться, а если убрать отрицательное число слева — увеличиться. Значит, по сравнению суммы с k нельзя однозначно решить, какую границу двигать.

2. Решение

Решение в лоб

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

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        answer = 0

        for left in range(len(nums)):
            current_sum = 0

            for right in range(left, len(nums)):
                current_sum += nums[right]

                if current_sum == k:
                    answer += 1

        return answer

Для каждой из N левых границ правая граница может пройти до конца массива. Поэтому такое решение работает за O(N^2) по времени и O(1) по дополнительной памяти.

Это полезный промежуточный шаг, но O(N^2) на собесе не проканает. Нужно избавиться от второго цикла.

Префиксная сумма — это накопленный итог

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

Для массива nums = [1, 2, 1, -1, 1] накопленные итоги выглядят так:

`num`	Накопленный итог
инициализация (затакт)	0
1	1
2	3
1	4
-1	3
1	4

В строке с числом уже учтено это число. Например, напротив второго числа 1 стоит итог 4, потому что мы уже сложили 1 + 2 + 1.

Нулевую сумму перед первым элементом можно воспринимать как затакт. Основная музыка — проход по массиву — ещё не началась, но мы уже подготовились к первому шагу. Часто этот шаг называют инициализацией: до запуска цикла мы записываем, что накопленный итог равен нулю.

Префиксные суммы можно считать и по-другому. Можно поставить 0 напротив первого элемента и считать, что в каждой позиции записана сумма до этого элемента. Тогда итог после последнего числа придётся вынести в дополнительную позицию за пределами массива. Это те же самые накопленные итоги, только сдвинутые на одну позицию. В этом разборе используем вариант с затактом: 0 стоит перед массивом, а каждая следующая строка уже включает текущий num.

Как получить сумму подмассива

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

Например, возьмём подмассив [2, 1], который начинается на индексе 1 и заканчивается на индексе 2:

накопленный итог в его конце равен 4;
накопленный итог перед его началом равен 1;
вычитаем: 4 - 1 = 3.

Мы вычли всё, что находилось до подмассива, и оставили только сумму его элементов.

Если подмассив начинается с самого начала массива, перед ним находится наш затакт с накопленным итогом 0. Например, для первых двух чисел [1, 2]:

3 - 0 = 3

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

Связь с Two Sum

По сути, это тот же паттерн, что и в Two Sum: идём слева направо, для текущего значения вычисляем нужный «комплемент» и проверяем, встречался ли он среди предыдущих значений.

В Two Sum уравнение выглядит так:

previous + current = target
previous = target - current

Здесь вместо отдельных элементов мы работаем с накопленными итогами, а сумма подмассива получается их разностью:

current_sum - previous_sum = k
previous_sum = current_sum - k

Значит, для текущего накопленного итога current_sum комплемент равен current_sum - k. Мы ищем, сколько раз такой итог уже встречался раньше.

Например, текущий итог равен 4, а k = 3. Ищем:

4 - 3 = 1

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

Есть только одно важное отличие от классической Two Sum. Там обычно достаточно либо узнать, существовал ли комплемент, либо найти одну пару индексов. Здесь нужно посчитать все подмассивы, поэтому мы считаем, сколько раз в прошлом встречался каждый предыдущий накопленный итог (один и тот же накопленный итог может встречаться несколько раз из-за нулей и других фрагментов с нулевой суммой).

Пройдём пример по шагам

Возьмём nums = [1, 2, 1, -1, 1] и k = 3. До начала цикла нам уже известен накопленный итог 0 — это затакт, который мы записали при инициализации.

`num`	Текущий итог	Ищем `итог - k`	Все итоги, которые видели раньше	Сколько раз видели нужный итог в прошлом	Новые подмассивы	Ответ
затакт (инициализация)	0	—	`[]`	—	—	0
1	1	-2	`[0]`	0	—	0
2	3	0	`[0, 1]`	1	`[1, 2]`	1
1	4	1	`[0, 1, 3]`	1	`[2, 1]`	2
-1	3	0	`[0, 1, 3, 4]`	1	`[1, 2, 1, -1]`	3
1	4	1	`[0, 1, 3, 4, 3]`	1	`[2, 1, -1, 1]`	4

Разберём, откуда взялся столбец «Сколько раз видели нужный итог в прошлом»:

В затакте мы ещё не взяли ни одного числа. Накопленный итог равен 0, и именно его сохраняем перед началом цикла.
Для первого num = 1 текущий итог равен 1. Ищем 1 - 3 = -2, но раньше видели только 0, поэтому прибавляем 0.
После num = 2 текущий итог равен 3. Ищем 3 - 3 = 0. Ноль встречался один раз — в затакте, поэтому засчитываем подмассив [1, 2].
После следующего num = 1 текущий итог равен 4. Ищем 4 - 3 = 1. Итог 1 встречался один раз, поэтому засчитываем подмассив [2, 1].
После num = -1 текущий итог снова равен 3. Ищем 0, который по-прежнему встречался один раз. Засчитываем подмассив [1, 2, 1, -1].
После последнего num = 1 текущий итог снова равен 4. Ищем 1, который встречался один раз. Засчитываем подмассив [2, 1, -1, 1].

Почему здесь удобен `Counter`

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

В таблице выше мы каждый раз просматривали все прошлые итоги и считали нужное значение. В коде такой просмотр вернул бы второй цикл и снова привёл бы к O(N^2).

Вместо списка используем Counter — готовую хеш-таблицу частот. Он хранит:

prefix_count[сумма] = сколько раз эта сумма уже встречалась

Теперь количество найденных комплементов можно получить одной операцией:

answer += prefix_count[prefix_sum - k]

Если такого итога в Counter нет, он вернёт 0. Если нужный накопленный итог встречался три раза, текущее число заканчивает сразу три разных подмассива с суммой k, и мы прибавим к ответу 3.

Зачем начинать с `Counter({0: 1})`

До первого элемента накопленный итог равен нулю. Это наш затакт, или инициализация перед циклом, поэтому заранее записываем:

prefix_count = Counter({0: 1})

Без этой записи мы потеряем все подходящие подмассивы, начинающиеся с индекса 0.

Например, для nums = [3] и k = 3 текущая сумма равна 3, а нужная предыдущая сумма — 3 - 3 = 0. Именно заранее добавленный ноль позволяет посчитать подмассив [3].

Алгоритм по шагам

Создаём Counter частот prefix_count = Counter({0: 1}).
Кладём в prefix_sum текущий накопленный итог, сначала 0.
Идём слева направо по nums и добавляем очередное число в prefix_sum.
Находим prefix_sum - k.
Прибавляем к ответу, сколько раз такой накопленный итог встречался раньше.
Увеличиваем частоту текущей prefix_sum в Counter.
После обработки всего массива возвращаем ответ.

Порядок последних двух действий важен: сначала считаем подходящие предыдущие суммы, и только потом записываем текущую. Иначе при k = 0 текущий накопленный итог ошибочно составит пару сам с собой, то есть мы посчитаем пустой подмассив.

Код на Python

from collections import Counter
from typing import List


class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        # До начала массива префиксная сумма 0 встретилась один раз
        prefix_count = Counter({0: 1})

        prefix_sum = 0
        answer = 0

        for num in nums:
            prefix_sum += num

            # Каждая предыдущая сумма prefix_sum - k
            # задаёт один подмассив с суммой k
            answer += prefix_count[prefix_sum - k]

            # Текущую сумму записываем только после подсчёта ответа
            prefix_count[prefix_sum] += 1

        return answer

Почему алгоритм работает

Рассмотрим любой подмассив с суммой k. Возьмём два накопленных итога:

текущий — в конце этого подмассива;
предыдущий — перед его началом.

Разность этих итогов равна сумме подмассива:

current_sum - previous_sum = k

Значит, нужный предыдущий итог равен current_sum - k. К моменту, когда цикл дошёл до конца подмассива, этот предыдущий итог уже записан в словаре. Поэтому алгоритм обязательно найдёт и посчитает подмассив.

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

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

Таким образом, алгоритм считает каждый подходящий подмассив ровно один раз — когда доходит до его последнего элемента.

Почему не подходит скользящее окно

В задачах с положительными числами часто можно поддерживать окно двумя указателями:

сумма слишком мала — расширяем окно;
сумма слишком велика — сужаем окно.

Здесь такой логики нет. Рассмотрим:

nums = [2, -1, 2]
k = 3

После первого элемента сумма равна 2. Добавляем следующий элемент и получаем 1: расширение окна уменьшило сумму. После третьего элемента сумма становится 3.

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

Краевые случаи

Подмассив начинается с первого элемента. Его учитывает начальная запись {0: 1}.
k = 0. Сначала читаем частоту prefix_sum - k, затем обновляем частоту текущей суммы — иначе посчитаем пустой подмассив.
В массиве есть отрицательные числа. Они не мешают формуле разности префиксных сумм.
В массиве есть нули. Одна префиксная сумма может повториться много раз, поэтому храним частоты.
Подходит один элемент. Он считается так же, как любой другой подмассив.
Подходящих подмассивов нет. Все обращения к отсутствующим ключам возвращают 0, и ответ остаётся нулевым.
Все числа равны нулю и k = 0. Подходят все подмассивы; повторяющиеся нулевые префиксные суммы позволяют корректно посчитать каждый из них.

3. Сложность алгоритма

Мы один раз проходим по массиву. На каждом шаге выполняем несколько операций со словарём, которые в среднем занимают O(1).

Итог по времени: O(N).

В худшем случае все префиксные суммы различны, и в словаре будет N + 1 ключей.

Итог по дополнительной памяти: O(N).

Главная идея задачи: это Two Sum на префиксных суммах. Для текущей суммы мы ищем среди прошлых префиксов комплемент current_prefix - k, а его частота показывает, сколько подмассивов с суммой k заканчивается в текущей позиции.

🐆