blob: 431cae5a1ce6c3eb74c82c21e8d9b035e772a5b7 (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
|
# 560. Subarray Sum Equals K
"""
given an array of integers 'nums' and an integer 'k', return the total number
of subarrays whose sum equals to 'k'. a subarray is a contiguous non-empty
sequence of elements within an array.
"""
class Solution(object):
def subarraySum(self, nums, k):
ans = 0
prefix_sum = 0
d = {0: 1}
for n in nums:
prefix_sum += n
if prefix_sum - k in d:
ans += d[prefix_sum - k]
if prefix_sum not in d:
d[prefix_sum] = 1
else:
d[prefix_sum] += 1
return ans
if __name__ == "__main__":
obj = Solution()
print(obj.subarraySum([1, 1, 1], 2)) # expect: 2
print(obj.subarraySum([1, 2, 3], 3)) # expect: 2
|