Leetcode 525. 连续数组（前缀和）

题目

给定一个二进制数组 nums , 找到含有相同数量的 0 和 1 的最长连续子数组，并返回该子数组的长度。

示例 1:

输入: nums = [0,1]
输出: 2
说明: [0, 1] 是具有相同数量0和1的最长连续子数组。

示例 2:

输入: nums = [0,1,0]
输出: 2
说明: [0, 1] (或 [1, 0]) 是具有相同数量0和1的最长连续子数组。

提示：

• 1 <= nums.length <= 10^5
• nums[i] 不是 0 就是 1

思路

先看数据量为10的5次方，所以肯定不能暴力做，那么可以借助前缀和来解决这个问题；
我们令prefix[i]代表序列[0,i]中剩余数字1的个数，注意这里如果遇到零则会抵消一个1；举个例子，对于序列[0,1,1,0]，他的prefix[3] 等于零，因为0和1的个数正好抵消了；
依次计算每个前缀和数组，如果发现前缀和之前出现过，说明存在序列包含相同数量的0和1，比较 之后更新答案；
如果发现前缀和之前没有出现过，则需要将它放入map中，key值为此时的前缀和，value值为此时的数组下标；

代码

class Solution {
    public int findMaxLength(int[] nums) {
        int res = 0;
        //map的key代表前缀和数组中数字1的个数，value代表相应key第一次出现的右端点下标
        Map<Integer, Integer> map = new HashMap<Integer, Integer>();
        int count = 0;
        //把-1放进去的原因是，如果子序列[0,2]包含0和1的数量相等，那么这个子序列的长度为3
        map.put(count, -1);
        for(int i = 0; i < nums.length; i ++) {
            int num = nums[i];
            if(num == 1) count++;
            else count--;
            if(map.containsKey(count)) {
                int startIndex = map.get(count);
                //startIndex是最早出现的下标，此时一定也是最长的
                res = Math.max(res, i - startIndex);
            } else {
                map.put(count, i);
            }
        }
        return res;
    }
}