pylist

560. Subarray Sum Equals K

Given an array of integers and an integer k, you need to find the total number of continuous subarrays whose sum equals to k.

Example 1:

Input:nums = [1,1,1], k = 2
Output: 2

Note:

Solution

Refined submission, use counter

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        counter = 0
        inv_cum = InvCumsum()
        for i in range(len(nums)):
            # k itself is a satisfied subarray
            if k == nums[i]:
                counter += 1
            # k and previous elements can construct satisfied subarrays
            counter += inv_cum.n_indexes(k - nums[i])
            # update
            inv_cum.append(nums[i])
        return counter
    
class InvCumsum:
    def __init__(self):
        self.sum_of_all = 0
        self.cumsum = dict() # without current element
        
    def append(self, x):
        self.cumsum.setdefault(self.sum_of_all, 0)
        self.cumsum[self.sum_of_all] += 1
        self.sum_of_all += x
        
    def n_indexes(self, inv_cum):
        cum = self.sum_of_all - inv_cum
        return self.cumsum.get(cum, 0)

Initial submission, use dict of list

class Solution:
    def subarraySum(self, nums: List[int], k: int) -> int:
        counter = 0
        inv_cum = inv_cumsum()
        # TODO: look init later
        for i in range(len(nums)):
            counter += len(inv_cum.indexes(k - nums[i]))
            if (k - nums[i]) == 0:
                counter += 1
            # update
            inv_cum.append(nums[i])
        return counter
    
class inv_cumsum:
    def __init__(self):
        self.sum_of_all = 0
        self.cumsum = dict() # without current element
        self.i = 0
        
    def append(self, x):
        self.cumsum.setdefault(self.sum_of_all, []).append(self.i)
        self.sum_of_all += x
        self.i += 1
        
    def indexes(self, inv_cum):
        cum = self.sum_of_all - inv_cum
        return self.cumsum.get(cum, [])

Care