aboutsummaryrefslogtreecommitdiff
path: root/subarr_sum_eq_k.py
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