Want Help Cracking FAANG?

(Then click this)

×
Back to Question Bank

2031. Count Subarrays With More Ones Than Zeros - Leetcode Solution

Code Implementation

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;
};
      

Problem Description

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:

  • Each subarray considered must have more 1s than 0s.
  • Elements can only be counted once per subarray; do not reuse elements across subarrays.
  • Return the total count of such subarrays.

Thought Process

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.

Solution Approach

To solve the problem efficiently, we use the following approach:

  1. Transform the Array: Replace every 1 with +1 and every 0 with -1. This way, the sum of a subarray is positive if and only if there are more 1s than 0s in that subarray.
  2. Prefix Sum (Balance): Compute a running prefix sum of this transformed array. For each index, the prefix sum represents the net balance of 1s over 0s up to that point.
  3. Key Observation: For a subarray ending at index 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].
  4. Counting Efficiently: As we iterate through the array, for each prefix sum (balance), we want to count how many previous prefix sums are less than the current balance. This tells us how many subarrays ending at the current index have more 1s than 0s.
  5. Implementation:
    • We can keep a map or array to count occurrences of each balance seen so far.
    • For each new balance, sum up the counts of all previous balances less than the current one.
    • Update the count for the current balance.
This approach avoids checking every possible subarray explicitly and leverages prefix sums and counting for efficiency.

Example Walkthrough

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):

  • Index 0: 0 (empty prefix)
  • Index 1: 1
  • Index 2: 0
  • Index 3: 1
  • Index 4: 0
  • Index 5: 1

Step 3: Count Subarrays
For each prefix sum, count how many previous prefix sums are less than the current one:
  • At balance 1 (index 1): previous balances = [0]; 0 < 1 → 1 subarray
  • At balance 0 (index 2): previous balances = [0,1]; none < 0 → 0
  • At balance 1 (index 3): previous balances = [0,1,0]; 0 < 1 (twice) → 2 subarrays
  • At balance 0 (index 4): previous balances = [0,1,0,1]; none < 0 → 0
  • At balance 1 (index 5): previous balances = [0,1,0,1,0]; 0 < 1 (three times) → 3 subarrays

Result: Total subarrays = 1 + 0 + 2 + 0 + 3 = 6
Thus, there are 6 subarrays with more 1s than 0s.

Time and Space Complexity

Brute-force approach:

  • Check all possible subarrays: O(n2) time, O(1) space.
  • This is too slow for large arrays.
Optimized approach (prefix sum + map):
  • For each index, we scan through the map of balances. In the worst case, this can be O(n) per index, leading to O(n2) time.
  • With further optimization (using a Binary Indexed Tree or ordered map), we can reduce this to O(n log n) time.
  • Space complexity: O(n) for storing prefix balances.
Summary: The provided code is simple and correct, but for best performance on large inputs, a more advanced data structure may be needed for O(n log n).

Summary

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.