subarray sum equals k

🏠

The subarray sum equals k problem asks:

Given an array of integers and an integer k, you need to find the total number of continuous subarrays whose sum equals to k.

Think about it this way: you're given a list of 'hops' you've taken down a road: 1, 2, 2, 0, 3, 2, 5. Initially you hopped 1 meter, then 2, and 2 again, then 0 meters etc.

You are then asked: how many of these hops travelled 5 meters consecutively? The answer is 6:

1              _
2          ___
3        _____
4      _____     
5 ________
6 ______
7 [1,2,2,0,3,2,5]

Rather than working with individual hop data, let's think of the total distance we travelled.

Work with accumulation values

So our total distance travelled at each hop is:

1 nums:          1,  2,  2,  0,  3,  2,  5
2 accumulation:  1,  3,  5,  5,  8,  10, 15

This is something that trips me up often with DP questions: you can often end up working in a totally different solution space. We are no longer concerned with our input values.. only their accumulation values. That can take quite a leap of the imagination (no pun intended).

Could i hop back 'k'?

Let's add our starting point 0 to the accumulator, think of it as the starting point. This 'total distance travelled' data is invaluable, because we can ask it questions, namely:

1 0, 1, 3, 5, 5, 8, 10, 15
2 -
3 |
4 |
5 Q: have i travelled 5 meters yet
6 A: nope
1 0, 1, 3, 5, 5, 8, 10, 15
2    -
3    |
4    |
5 Q: have i travelled 5 meters yet
6 A: nope
1 0, 1, 3, 5, 5, 8, 10, 15
2       -
3       |
4       |
5 Q: have i travelled 5 meters yet
6 A: nope
1 0, 1, 3, 5, 5, 8, 10, 15
2          -
3          |
4          |
5 Q: have i travelled 5 meters yet
6 A: ah yes you have

Let's think of another question to ask:

1 0, 1, 3, 5, 5, 8, 10, 15
2 x........-
3          |
4          |
5 Q: if i hop back 5 meters, do i land on a hop
6 A: yes

acc - k

In other words, the question we are asking is: is there an accumulation value equal to acc - k.

Let's count the 'hits':

1 0, 1, 3, 5, 5, 8, 10, 15
2 x........x
3 x...........x
4       x........x
5          x..x......x
6                x......x

Now it gets a little trickier. When we travel 10 meters, and ask do i land on a hop if i travel 5 meters back? the answer is: actually you could land on two hops. Since there is a 0 in the array, there are two subarrays that add up to the same accumulation value. We need to account for that.

An easy to understand implementation

1 from itertools import accumulate
2 
3 def subarraySum(nums, k):
4     total = 0
5     accum = [0] + [*accumulate(nums)]
6     for i, acc in enumerate(accum[1:]):
7         total += accum[: i + 1].count(acc - k)
8     return total

We begin by adding out starting position 0 to our accumulation array (when we hop back, it's valid to land on the starting position). We then iterate over our accumulate (barring 0).

We then ask if i hop back k meters, how many valid hops do i find?. This works, but is highly inefficient.

A more efficient implementation

1 from collections import Counter
2 from itertools import accumulate
3 
4 def subarraySum(nums, k):
5     count, total = Counter({0: 1}), 0
6     for acc in accumulate(nums):
7         total += count[acc - k]
8         count[acc] += 1
9     return total

A more efficient implementation is to count the acc (distance travelled) values in a counter, on line 8 as we go along (as we iterate over accumulate). This is purely for the O(1) lookup.

We initialise the counter with {0:1} since if we land on our starting position, we've made a valid hop.

We'll want to perform the lookup (ask the question, on line 7) before adding our accumulation value (distance) to the counter.