Subarrays Summing to K

Hard · rating 1600 · Arrays, Prefix Sum, Hashing

The first line contains n and a target k. The second line contains n integers. Print how many contiguous subarrays sum to exactly k.

Constraints: 1 ≤ n ≤ 105, |ai| ≤ 104

Editorial

Approach

Let prefix[i] be the sum of the first i elements. A subarray (i, j] sums to k exactly when prefix[j] − prefix[i] = k. So as you sweep the running prefix, count how many earlier prefixes equal run − k using a hash map seeded with {0: 1} for subarrays starting at index 0.

from collections import Counter
seen = Counter({0: 1})
run = count = 0
for x in a:
    run += x
    count += seen[run - k]
    seen[run] += 1
print(count)

Why the seed

The {0: 1} entry represents the empty prefix, letting subarrays that begin at the very start be counted.

Complexity

Time: O(n). Space: O(n).

Related problems

Open Subarrays Summing to K in Code Arena →