Need help to understand logic in the code

May 5, 2022

The below code is returning the total number of subarrays whose sum equals to k.
I am trying to understand logic in the below code:

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

dic = {0:1}
res = 0
curr = 0
for i in range(len(nums)):
        curr += nums[i]
        res += dic.get(curr-k, 0)  #need help to understand the logic
        dic[curr] = dic.get(curr, 0) + 1  #need help to understand the logic 
print(dic)
print(res)

Output:

{0: 1, 1: 1, 4: 1, 5: 1}

>Solution :

This is code that uses the strategy of prefix sums.

Prefix sums are the sums of all previous numbers in the array at a certain index. If you do prefix[f] - prefix[i], you get the sum of the array between f and i (f > i).

It has a dictionary that keeps track of prefix sums to find the desired k value.

nums = [1,3,1] #input
k = 4 #desired value

dic = {0:1} # empty set = 0
res = 0 #amount of subarrays with sum k found
curr = 0 #running total
for i in range(len(nums)):
        curr += nums[i]
        res += dic.get(curr-k, 0)  # gets the number of previous prefix sums 
            #that can be subtracted from the current to get k
        dic[curr] = dic.get(curr, 0) + 1  #puts the current prefix sum in the dictionary
print(dic)
print(res)

Because prefix[f] - prefix[i] is sum between f and i, if you have prefix[f], you can find the value of prefix[i] to get k. So then you just count the number of prefix[i]s you have found.

dict.get(key) is used instead of dict[key] because if the prefix value has not been found (yet), then it won’t cause a KeyError, and would return 0 (the actual count of the value) instead.

PS: If this is related to coding olympiads, there is this and this.