Prefix Sum

525. Contiguous Array

Approach:

1
class Solution:
2
    def findMaxLength(self, nums: List[int]) -> int:
3
        ans = 0
4
        memo = {}
5
        presum = [0] * len(nums)
6
        for i in range(len(nums)):
7
            if i > 0:
8
                presum[i] += presum[i-1]
9
            if nums[i] == 1:
10
                presum[i] += 1
11
            else:
12
                presum[i] += -1
13
            if presum[i] not in memo:
14
                memo[presum[i]] = i
15
            if presum[i] == 0:
16
                ans = max(ans, i+1)
17
            if presum[i] in memo:
18
                ans = max(ans, i - memo[presum[i]])
19
        return ans

Optimization:

As a matter of fact, we don’t need to use an array to store the prefix sum. Note that in the previous solution, only presum[i] is used in each iteration, so we can simply use a variable to store the latest prefix sum.
There is also another trick to simplify the code. Notice that the if presum[i] == 0 was used to check the current prefix sum is 0. This particaular case can be merged by initializing memo with memo = {0: -1}. In this case presum[i] == 0 is equivalent to presum[i] in memo and i+1 is identical to i - memo[presum[i]].

The final version

1
class Solution:
2
    def findMaxLength(self, nums: List[int]) -> int:
3
        ans = 0
4
        memo = {0: -1}
5
        total = 0
6
        for i in range(len(nums)):
7
            if nums[i] == 1:
8
                total += 1
9
            else:
10
                total += -1
11
            if total not in memo:
12
                memo[total] = i
13
            ans = max(ans, i - memo[total])
14
        return ans

523. Continuous Subarray Sum

Leetcode

Example 1:

Input: nums = [23,2,4,6,7], k = 6
Output: true
Explanation: [2, 4] is a continuous subarray of size 2 whose elements sum up to 6.

Example 2:

Input: nums = [23,2,6,4,7], k = 6
Output: true
Explanation: [23, 2, 6, 4, 7] is an continuous subarray of size 5 whose elements sum up to 42.
42 is a multiple of 6 because 42 = 7 * 6 and 7 is an integer.

Example 3:

Input: nums = [23,2,6,4,7], k = 13
Output: false

Approach:

1
class Solution:
2
    def checkSubarraySum(self, nums: List[int], k: int) -> bool:
3
        total = 0
4
        memo = {0: -1}
5
        for i in range(len(nums)):
6
            total += nums[i]
7
            total %= k
8
            if total not in memo:
9
                memo[total] = i
10
            elif i - memo[total] >= 2:
11
                return True
12
        return False