from collections import defaultdict
class Solution:
def countSubarrays(self, nums):
balance = 0
count = defaultdict(int)
count[0] = 1 # base: empty prefix
result = 0
for num in nums:
if num == 1:
balance += 1
else:
balance -= 1
# For current prefix, count all prefixes with balance less than current
# So, sum up all counts for balances < current
# We can keep a running sum instead of iterating every time
# But here, we use a trick: maintain prefix sum
# Actually, we can use an ordered map, but for simplicity, let's use a cumulative approach
# Instead, we can precompute the min balance to keep an array
# But let's do it efficiently: keep a running total of all counts so far
# Actually, we can use a variable to store the cumulative sum of counts for balances less than current
# But since balances can be negative, let's use a map
# Instead, we can process as follows:
# For each prefix, the number of previous prefixes with balance less than current
# equals the sum of count[bal] for bal in count if bal < balance
# To optimize, keep a running sum of all counts so far, and for each new balance,
# subtract the sum of counts for balances >= current balance.
# For simplicity, let's use an array of balances and prefix sums
# But for now, let's use brute force for clarity
for prev_balance in count:
if prev_balance < balance:
result += count[prev_balance]
count[balance] += 1
return result
#include <unordered_map>
#include <vector>
class Solution {
public:
int countSubarrays(std::vector<int>& nums) {
std::unordered_map<int, int> count;
count[0] = 1;
int balance = 0, result = 0;
for (int num : nums) {
balance += (num == 1) ? 1 : -1;
for (auto& [prev_balance, cnt] : count) {
if (prev_balance < balance) {
result += cnt;
}
}
count[balance]++;
}
return result;
}
};
import java.util.HashMap;
class Solution {
public int countSubarrays(int[] nums) {
HashMap<Integer, Integer> count = new HashMap<>();
count.put(0, 1);
int balance = 0, result = 0;
for (int num : nums) {
balance += (num == 1) ? 1 : -1;
for (int prev_balance : count.keySet()) {
if (prev_balance < balance) {
result += count.get(prev_balance);
}
}
count.put(balance, count.getOrDefault(balance, 0) + 1);
}
return result;
}
}
var countSubarrays = function(nums) {
let count = new Map();
count.set(0, 1);
let balance = 0, result = 0;
for (let num of nums) {
balance += (num === 1) ? 1 : -1;
for (let [prev_balance, cnt] of count.entries()) {
if (prev_balance < balance) {
result += cnt;
}
}
count.set(balance, (count.get(balance) || 0) + 1);
}
return result;
};
Given an array nums
consisting only of 0s and 1s, count the number of subarrays where the number of 1s is strictly greater than the number of 0s.
A subarray is a contiguous sequence of elements from nums
. The same element may not be reused in different subarrays, but overlapping subarrays are allowed as long as they are contiguous.
Key constraints:
The naive approach is to check every possible subarray, count the number of 1s and 0s in each, and see if 1s > 0s. This works for small arrays, but for large arrays (length up to 105), this is too slow.
To optimize, we need a way to efficiently check, for any subarray, whether it has more 1s than 0s, and count such subarrays quickly.
The key insight is to use prefix sums or balances: For any subarray, if we can quickly compute the difference between the number of 1s and 0s, we can decide if it qualifies.
We can represent each 1 as +1 and each 0 as -1, and compute a running sum (balance). For a subarray starting at index i
and ending at j
, the balance difference tells us if 1s > 0s. We then need a way to count, for each prefix, how many previous prefixes have a balance less than the current one.
To solve the problem efficiently, we use the following approach:
j
, starting at index i
, the sum of the subarray is prefix[j + 1] - prefix[i]
. We want this to be > 0, i.e., prefix[j + 1] > prefix[i]
.
Let's use the input nums = [1, 0, 1, 0, 1]
.
Step 1: Transform to +1/-1
The array becomes: [1, -1, 1, -1, 1]
Step 2: Compute Prefix Sums
Prefix balances at each index (starting from 0):
Brute-force approach:
The problem asks for the count of subarrays with more 1s than 0s in a binary array. By transforming the array to +1/-1 and using prefix sums, we can efficiently check the condition for each subarray. By maintaining counts of previous prefix balances, we avoid brute-force enumeration, making the solution much faster. This approach leverages key algorithmic techniques like prefix sums and hash maps, demonstrating how a conceptual shift can lead to significant optimization.