Need help to understand logic in the code-CodePudding

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}

CodePudding user response：

This is code that uses the strategy of prefix sums (or running totals, cumulative sums, scans, etc.).

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 an example problem, the solution, and more explanations on prefix sums.