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.

Recommend that you read this first: Algorithms to Find Maximum Size Subarray (Contiguous) Sum That Equals k. Then, the problem can be transformed to: Given an array of numbers that only contain +1 or -1, return the maximum size subarray (contiguous) that sum up to 0.

Possibly, this (Algorithms to Count Subarray (Contiguous) Sum That Equals k) is similar, but do the counting instead of returning the maximum subarray.

Then, we can use a hash map to store the prefix sums of the array. We store the prefix sums as keys, and the values are the first-met indices. Then, if prefix[sum1] = i, and prefix[sum2] = j, the sum of numbers from index i + 1 to j is sum2-sum1.

Therefore, if sum1 is equal to sum2, we can record the maximum subarray with length (j-i). Special case is when sum is zero thus we find a longest balanced pairs of zeros and ones.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

class Solution {
public:
    int findMaxLength(vector<int>& nums) {
        if (nums.empty()) return 0;
        if (nums.size() == 1) return 0;
        unordered_map<int, int> prefix;
        int sum = 0;
        int ans = 0;
        for (int i = 0; i < nums.size(); ++ i) {
            sum += nums[i] == 0 ? 1 : -1;
            if (sum == 0) {
                ans = i + 1;
            }
            if (prefix.find(sum) != prefix.end()) {
                ans = max(ans, i - prefix[sum]);
            } else {
                prefix[sum] = i;
            }
        }
        return ans;
    }
};

class Solution {
public:
    int findMaxLength(vector<int>& nums) {
        if (nums.empty()) return 0;
        if (nums.size() == 1) return 0;
        unordered_map<int, int> prefix;
        int sum = 0;
        int ans = 0;
        for (int i = 0; i < nums.size(); ++ i) {
            sum += nums[i] == 0 ? 1 : -1;
            if (sum == 0) {
                ans = i + 1;
            }
            if (prefix.find(sum) != prefix.end()) {
                ans = max(ans, i - prefix[sum]);
            } else {
                prefix[sum] = i;
            }
        }
        return ans;
    }
};

The space requirement is O(N) where we use a hash map i.e. unordered_map in C++. The time complexity is O(N) where we need to iterate once the entire array. We can set the prefix hash map with initial prefix[0] = -1 so that we can remove the special case handling inside the loop.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20

class Solution {
public:
    int findMaxLength(vector<int>& nums) {
        if (nums.empty()) return 0;
        if (nums.size() == 1) return 0;
        unordered_map<int, int> prefix;
        int sum = 0;
        prefix[0] = -1;  // special case
        int ans = 0;
        for (int i = 0; i < nums.size(); ++ i) {
            sum += nums[i] == 0 ? 1 : -1;
            if (prefix.find(sum) != prefix.end()) {
                ans = max(ans, i - prefix[sum]);
            } else {
                prefix[sum] = i;
            }
        }
        return ans;
    }
};

class Solution {
public:
    int findMaxLength(vector<int>& nums) {
        if (nums.empty()) return 0;
        if (nums.size() == 1) return 0;
        unordered_map<int, int> prefix;
        int sum = 0;
        prefix[0] = -1;  // special case
        int ans = 0;
        for (int i = 0; i < nums.size(); ++ i) {
            sum += nums[i] == 0 ? 1 : -1;
            if (prefix.find(sum) != prefix.end()) {
                ans = max(ans, i - prefix[sum]);
            } else {
                prefix[sum] = i;
            }
        }
        return ans;
    }
};

–EOF (The Ultimate Computing & Technology Blog) —