動機

用prefix sum紀錄偏移到哪邊了

Problem

Given a binary array nums, return the maximum length of a contiguous subarray with an equal number of 0 and 1.

 

Example 1:

Input: nums = [0,1]Output: 2Explanation: [0, 1] is the longest contiguous subarray with an equal number of 0 and 1.

Example 2:

Input: nums = [0,1,0]Output: 2Explanation: [0, 1] (or [1, 0]) is a longest contiguous subarray with equal number of 0 and 1.

 

Constraints:

    t
  • 1 <= nums.length <= 105
    • t
  • nums[i] is either 0 or 1.

Sol

如果總和是0，就是平衡的，如果有多，只要把多的扣掉就好，就是prefix sum

class Solution:
    def findMaxLength(self, nums: List[int]) -> int:
        l2i = {0:-1}
        cnt,ret = 0,0
        
        for (i,n) in enumerate([-1 if n == 0 else n for n in nums]):
            cnt += n
            if cnt in l2i:
                ret = max(i-l2i[cnt],ret)
            else:
                l2i[cnt] = i
        return ret