[leetcode] 525. Contiguous Array

Description

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

Example 1:

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

Example 2:

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

Note: The length of the given binary array will not exceed 50,000.

分析

题目的意思是：给你一个只有0或者1的数组，然后截取子数组使得0的数量跟1的数量是一样多的，求这种子数组的个数。

• 需要用到一个trick，遇到1就加1，遇到0，就减1，这样如果某个子数组和为0，就说明0和1的个数相等。
• 我们用一个哈希表建立子数组之和跟结尾位置的坐标之间的映射。如果某个子数组之和在哈希表里存在了，说明当前子数组减去哈希表中存的那个子数字，得到的结果是中间一段子数组之和，必然为0，说明0和1的个数相等，我们更新结果res

C++实现

class Solution {
public:
    int findMaxLength(vector<int>& nums) {
        int n=nums.size();
        int res=0;
        unordered_map<int,int> m{{0,-1}};
        int sum=0;
        for(int i=0;i<n;i++){
            sum+= nums[i]==1 ? 1:-1;
            if(m.count(sum)){
                res=max(res,i-m[sum]);
            }else{
                m[sum]=i;
            }
        }
        return res;
    }
};

Python实现

使用一个字典来跟踪每个位置的前缀和，并记录每个前缀和第一次出现的位置。当找到相同的前缀和时，说明这两个位置之间的子数组包含相同数量的0和1。然后计算它们的距离，即子数组的长度。

class Solution:
    def findMaxLength(self, nums: List[int]) -> int:
    	# 初始化字典，key为前缀和，value为第一次出现该前缀和的位置
        prefix_sum = {0:-1}
        res = 0
        cnt = 0
        for i in range(len(nums)):
            if nums[i]==0:
                cnt-=1
            else:
                cnt+=1
            if cnt in prefix_sum:
                res = max(res, i - prefix_sum[cnt])
            else:
                prefix_sum[cnt]=i
        return res

参考文献

[LeetCode] Contiguous Array 邻近数组