Subarray Sum Equals K

+10

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

Функция

subarraySum(nums: integer-array, k: integer) → integer

numsinteger-array: массив целых чисел, который может содержать отрицательные значения и нули
kinteger: сумма, которой должен достичь подмассив, чтобы его засчитали
Возвращаетinteger: количество подмассивов, сумма элементов которых равна k

Ограничения

1 ≤ nums.length ≤ 2 × 10⁴
-1000 ≤ nums[i] ≤ 1000
-10⁷ ≤ k ≤ 10⁷
У массива такой длины не более 200,010,000 подмассивов, поэтому ответ помещается в 32-разрядное знаковое целое число.

Примеры

Ввод: nums = [3, 4, -7, 1, 3, 3, 1, -4]k = 7
Вывод: 4
Пояснение: Четыре последовательности в сумме дают 7: [3, 4], [1, 3, 3], [3, 3, 1] и [3, 4, -7, 1, 3, 3]. В последней последовательности -7 компенсирует 3 и 4, и позднее сумма снова достигает 7, поэтому последовательность может подойти, даже если её сумма уже превысила k.

Ввод: nums = [1, -1, 0]k = 0
Вывод: 3
Пояснение: Три подмассива дают в сумме 0: [1, -1], [0] и весь массив [1, -1, 0]. Последовательность [-1, 0] даёт в сумме -1, поэтому не учитывается.

Ввод: nums = [2, 2, 2]k = 4
Вывод: 2
Пояснение: Последовательность [2, 2] с индексами 0 и 1 и последовательность [2, 2] с индексами 1 и 2 содержат одинаковые значения, но находятся в разных позициях, поэтому обе учитываются. Сумма всего массива равна 6.

Дополнительный вопрос

Как изменить решение, чтобы оно возвращало длину самого длинного подмассива, сумма элементов которого равна k, по-прежнему за время O(n)?

Подсказки

Открывайте по одной. Каждая подсказывает чуть больше.

Открыто 0 из 3

Решение

Массив из n чисел имеет n(n+1)/2 подмассивов — около 2 × 10^8 при n = 2 × 10^4, поэтому складывать числа в каждом из них слишком медленно. Отрицательные значения также исключают использование скользящего окна: сумма окна может уменьшаться, а затем снова увеличиваться, поэтому нет правила, которое подскажет, когда его сужать. Ключевая идея решения — представить сумму каждого подмассива как разность двух префиксных сумм. Чтобы посчитать подмассивы, заканчивающиеся на текущем элементе и дающие в сумме k, нужно посчитать, сколько предыдущих префиксных сумм равны текущей сумме минус k; хеш-таблица позволяет сделать это за один проход.

Каждое начало с накопительным итогом

Верно, но не успевает на самых больших тестах

Идея

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

Вам не нужен третий цикл, чтобы складывать элементы каждого подмассива. Зафиксируйте start, затем перемещайте end на одну позицию вправо и добавляйте nums[end] к текущему total. В total всегда хранится сумма элементов от start до end, поэтому для каждого подмассива требуется одно сложение и одно сравнение.

Не останавливайтесь, когда сумма достигает k или превышает его. Более позднее отрицательное значение может снова уменьшить её: в первом примере сумма от индекса 0 равна 3, 7, 0, 1, 4, 7, поэтому для этого начального индекса второе совпадение находится на индексе 5.

Вычислительная стоимость определяется количеством пар. При n = 2 × 10^4 их примерно 2 × 10^8 — это подходит для C, но слишком медленно для Python, Ruby или R.

Алгоритм

Установи count равным 0.
Для каждого start от 0 до n-1 установи total равным 0.
Для каждого end от start до n-1 прибавляй nums[end] к total.
Если total равен k, прибавь 1 к count и в любом случае продолжай.
Верни count.

def subarraySum(nums, k):
    n = len(nums)
    count = 0
    for start in range(n):
        total = 0
        # Grow the subarray that begins at start, one element at a time
        for end in range(start, n):
            total += nums[end]
            if total == k:
                count += 1
    return count

Префиксные суммы с картой подсчётов

Идея

Пусть prefix[j] — сумма первых j элементов, где prefix[0] = 0 для пустого префикса. Сумма подмассива от индекса i до индекса j-1 равна prefix[j] - prefix[i]. Поэтому сумма подмассива, заканчивающегося на текущем элементе, равна k тогда и только тогда, когда более ранняя сумма префикса равна текущей сумме префикса минус k. Каждая такая более ранняя сумма префикса указывает, где начинается один подходящий подмассив.

Пройди по массиву один раз. Поддерживай текущую сумму префикса и хеш-таблицу seen, в которой для каждой суммы префикса хранится, сколько раз она встречалась. На каждом элементе сначала прибавь seen[prefix - k] к счётчику, а затем запиши текущую сумму префикса. Поиск перед записью не позволяет подмассиву быть пустым: при k = 0 запись сначала привела бы к совпадению текущей суммы префикса с самой собой.

Рассмотрим первый пример с k = 7. Суммы префиксов равны 0, 3, 7, 0, 1, 4, 7, 8, 4. Когда после индекса 1 сумма префикса достигает 7, в таблице хранится один 0, который даёт [3, 4]. Когда после индекса 5 она снова достигает 7, в таблице хранится два нуля — пустой префикс и префикс после -7; вместе они дают [3, 4, -7, 1, 3, 3] и [1, 3, 3]. При значении 8 после индекса 6 в таблице хранится одна 1, которая даёт [3, 3, 1]. Итого получается 4.

Именно начальное значение 0, учтённое в таблице один раз, позволяет посчитать подмассивы, начинающиеся с индекса 0. Важно использовать таблицу счётчиков, а не множество, потому что одна и та же сумма префикса может повторяться, и каждая её копия соответствует началу другого подмассива. Для каждого элемента выполняются один поиск и одно обновление, поэтому время работы составляет O(n), а таблица содержит не более n+1 ключей.

Алгоритм

Создай карту seen с seen[0] = 1 и установи для prefix и count значение 0.
Для каждого элемента добавь его к prefix.
Добавь seen[prefix - k] к count, считая отсутствующий ключ равным 0.
Добавь 1 к seen[prefix].
Верни count.

def subarraySum(nums, k):
    # seen[p] = how many prefixes so far add up to p; the empty prefix adds up to 0
    seen = {0: 1}
    prefix = 0
    count = 0
    for num in nums:
        prefix += num
        # Each earlier prefix equal to prefix - k starts a subarray that ends here and sums to k
        count += seen.get(prefix - k, 0)
        seen[prefix] = seen.get(prefix, 0) + 1
    return count

Ловушки и крайние случаи

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

Скользящее окно, которое сужается, как только сумма превышает k, не работает с отрицательными значениями. В первом примере оно возвращает 2 вместо 4: окно сохраняет левую границу на индексе 0, пока сумма не превысит 7 на индексе 6, поэтому оно ни разу не проверяет [1, 3, 3] или [3, 3, 1].
Если не задать seen[0] = 1, будут пропущены все подмассивы, начинающиеся с индекса 0. Для nums = [5] и k = 5 алгоритм возвращает 0 вместо 1.
Если записывать текущую префиксную сумму до поиска, при k, равном 0, будут учитываться пустые подмассивы. Для [1, -1, 0] алгоритм возвращает 6 вместо 3.
Использование множества префиксных сумм вместо карты счётчиков приводит к тому, что повторения учитываются не полностью. Для [0, 0, 0] и k = 0 ответ равен 6, потому что каждая предыдущая копия той же префиксной суммы задаёт начало отдельного подмассива.
В методе полного перебора прерывать внутренний цикл, когда сумма превысит k, неправильно по той же причине, что и в случае со скользящим окном.

Частые вопросы3

Какова временная сложность задачи «Subarray Sum Equals K»?

Решение с префиксной суммой и хеш-таблицей выполняется за время O(n) и требует O(n) дополнительной памяти: один проход, один поиск и одно обновление для каждого элемента. Проверка каждого подмассива с текущей суммой занимает время O(n²), а суммирование каждого подмассива с нуля — O(n³).

Почему метод скользящего окна не работает для задачи «Сумма подмассива равна K»?

Метод скользящего окна опирается на то, что сумма увеличивается при расширении окна и уменьшается при его сужении; это верно только тогда, когда все значения положительные. При наличии отрицательных значений окно, сумма которого уже слишком велика, всё ещё может дать нужный результат после дальнейшего расширения, поэтому нет правила, подсказывающего, когда перемещать левую границу. Если бы все значения были положительными, метод скользящего окна решил бы эту задачу за время O(n) и с использованием O(1) памяти.

Почему в хеш-таблице ключу 0 соответствует значение 1?

Эта запись обозначает пустой префикс перед первым элементом, сумма которого равна 0. Сумма подмассива, начинающегося с индекса 0, равна текущей сумме префикса минус сумма этого пустого префикса, поэтому без этой записи такие подмассивы никогда не учитываются. Для nums = [5] и k = 5 поиск 5 - 5 = 0 находит эту запись и возвращает 1.

def subarraySum(nums, k):
    # Напишите код здесь

Тест #1

Тест #2

Тест #3

Аргументы

nums = [3, 4, -7, 1, 3, 3, 1, -4]
k = 7

Ожидаемый вывод