diff options
Diffstat (limited to 'subarr_sum_eq_k.py')
-rw-r--r-- | subarr_sum_eq_k.py | 29 |
1 files changed, 29 insertions, 0 deletions
diff --git a/subarr_sum_eq_k.py b/subarr_sum_eq_k.py new file mode 100644 index 0000000..431cae5 --- /dev/null +++ b/subarr_sum_eq_k.py @@ -0,0 +1,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 |