Welcome to Subscribe On Youtube
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 1
s minus count of 0
s 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.