525.Contiguous Array

Tags: Medium Hash Table

Links: https://leetcode.com/problems/contiguous-array/

---

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.

---

class Solution {
public:
    int findMaxLength(vector<int>& nums) {
        std::ios_base::sync_with_stdio(false);
        cin.tie(NULL);
        cout.tie(NULL);

        int n = nums.size();
        for (int i = 0; i < n; ++i) {
            if (!nums[i]) nums[i] = -1;
        }

        vector<int> preSum(n + 1, 0);
        for (int i = 1; i <= n; ++i) {
            preSum[i] = preSum[i - 1] + nums[i - 1];
        }

        unordered_map<int, vector<int>> um;
        for (int i = 0; i <= n; ++i) {
            um[preSum[i]].push_back(i);
        }

        int maxLen = -1;
        for (auto & e : um) {
            if (e.second.size() >= 2) {
                int length = e.second.size();
                maxLen = max(maxLen, e.second[length - 1] - e.second[0]);
            }
        }

        return maxLen == -1 ? 0 : maxLen;
    }
};

先把序列里的0全都换成-1，然后求前缀和，比如序列

0 0 0 1 0 1 1 1
变为
-1 -1 -1 1 -1 1 1 1
前缀和
0 -1 -2 -3 -2 -3 -2 -1 0

发现如果两个数前缀和相同，那么下标相减就是最大长度，那么从首元素开始寻找，找到就直接返回。这样看似不错，但是如果最坏情况，时间复杂度就是$O(n^2)$。所以可以考虑用哈希表来解决，用一个键为前缀和的大小，值是一个vector，存储所有前缀和等于键的下标，因为是按照顺序访问的，所以每个vector只需要关注第一个下标和最后一个下标。

但是上述代码还是有很多可以优化的地方。

class Solution {
public:
    int findMaxLength(vector<int>& nums) {
        std::ios_base::sync_with_stdio(false);
        cin.tie(NULL);
        cout.tie(NULL);

        int n = nums.size();
        int res = 0, sum = 0;
        unordered_map<int, int> um{{0, -1}};
        for (int i = 0; i < n; ++i) {
            sum += nums[i] == 1 ? 1 : -1;
            if (um.find(sum) != um.end()) {
                res = max(res, i - um[sum]);
            }
            else um[sum] = i;
        }

        return res;
    }
};

思路和初始思路是一致的。