Formatted question description: https://leetcode.ca/all/2031.html

2031. Count Subarrays With More Ones Than Zeros (Medium)

You are given a binary array nums containing only the integers 0 and 1. Return the number of subarrays in nums that have more 1's than 0's. Since the answer may be very large, return it modulo 109 + 7.

A subarray is a contiguous sequence of elements within an array.

 

Example 1:

Input: nums = [0,1,1,0,1]
Output: 9
Explanation:
The subarrays of size 1 that have more ones than zeros are: [1], [1], [1]
The subarrays of size 2 that have more ones than zeros are: [1,1]
The subarrays of size 3 that have more ones than zeros are: [0,1,1], [1,1,0], [1,0,1]
The subarrays of size 4 that have more ones than zeros are: [1,1,0,1]
The subarrays of size 5 that have more ones than zeros are: [0,1,1,0,1]

Example 2:

Input: nums = [0]
Output: 0
Explanation:
No subarrays have more ones than zeros.

Example 3:

Input: nums = [1]
Output: 1
Explanation:
The subarrays of size 1 that have more ones than zeros are: [1]

 

Constraints:

  • 1 <= nums.length <= 105
  • 0 <= nums[i] <= 1

Companies:
Google

Related Topics:
Array, Binary Search, Divide and Conquer, Binary Indexed Tree, Segment Tree, Merge Sort, Ordered Set

Similar Questions:

Solution 1.

index:     0  1  2  3  4
A:         0  1  1  0  1
Diff:   0 -1  0  1  0  1 // Count(1) - Count(0)
Count:     0 +1 +3 +1 +4 // Sum of number of diffs less than the current diff.

Let diff[i] be the count of 1s minus count of 0s before A[i] inclusive.

For each A[i], it will add “sum of number of diffs less than the current diff” to the answer.

So we need a data structure with which we can query a range sum. Segment tree and Binary Indexed Tree are good for this purpose.

This implementation uses BIT.

// OJ: https://leetcode.com/problems/count-subarrays-with-more-ones-than-zeros/
// Time: O(NlogN)
// Space: O(N)
// Ref: https://leetcode.com/problems/count-subarrays-with-more-ones-than-zeros/discuss/1512961/BIT-vs.-O(n)
const int N = 200000, mod = 1e9 + 7;
int bt[N + 1] = {};
class Solution {
    int sum(int i) {
        int ans = 0;
        for (++i; i > 0; i -= i & -i) ans += bt[i];
        return ans;
    }
    void update(int i, int val) {
        for (++i; i <= N; i += i & -i) bt[i] += val;
    }
public:
    int subarraysWithMoreZerosThanOnes(vector<int>& A) {
        int ans = 0, diff = 0;
        memset(bt, 0, sizeof(bt));
        update(N / 2, 1);
        for (int n : A) {
            diff += n ? 1 : -1;
            update(N / 2 + diff, 1);
            ans = (ans + sum(N / 2 + diff - 1)) % mod;
        }
        return ans;
    }
};

Solution 2. Divide and Conquer (Merge Sort)

// OJ: https://leetcode.com/problems/count-subarrays-with-more-ones-than-zeros/
// Time: O(NlogN)
// Space: O(N)
class Solution {
public:
    int subarraysWithMoreZerosThanOnes(vector<int>& A) {
        vector<int> diff(A.size() + 1), tmp(A.size() + 1);
        for (int i = 0, d = 0; i < A.size(); ++i) {
            d += A[i] == 1 ? 1 : -1;
            diff[i + 1] = d;
        }
        function<int(int, int)> mergeSort = [&](int begin, int end) -> int {
            if (begin + 1 >= end) return 0;
            long mid = (begin + end) / 2, mod = 1e9 + 7, ans = (mergeSort(begin, mid) + mergeSort(mid, end)) % mod;
            for (int i = begin, j = mid, k = begin; i < mid || j < end; ++k) {
                if (j == end || (i < mid && diff[i] < diff[j])) tmp[k] = diff[i++];
                else {
                    ans = (ans + i - begin) % mod;
                    tmp[k] = diff[j++];
                }
            }
            for (int i = begin; i < end; ++i) diff[i] = tmp[i];
            return ans;
        };
        return mergeSort(0, diff.size());
    }
};

Solution 3.

// OJ: https://leetcode.com/problems/count-subarrays-with-more-ones-than-zeros/
// Time: O(N)
// Space: O(N)
// Ref: https://leetcode.com/problems/count-subarrays-with-more-ones-than-zeros/discuss/1512961/BIT-vs.-O(n)
const int N = 200000, mod = 1e9 + 7;
int bt[N + 1] = {};
class Solution {
public:
    int subarraysWithMoreZerosThanOnes(vector<int>& A) {
        int ans = 0, diff = 0, cnt = 0;
        memset(bt, 0, sizeof(bt));
        bt[N / 2] = 1;
        for (int n : A) {
            diff += n ? 1 : -1;
            cnt += n ? bt[N / 2 + diff - 1] : -bt[N / 2 + diff];
            ans = (ans + cnt) % mod;
            ++bt[N / 2 + diff];
        }
        return ans;
    }
};

TODO

Add notes to Solution 1 and 3.

All Problems

All Solutions