Problem Description

You are given a binary array nums that contains only integers 0 and 1. Your task is to find the total number of subarrays where the count of 1s is greater than the count of 0s.

A subarray is defined as a contiguous sequence of elements within the array. For example, if nums = [1, 0, 1, 1], then [1, 0, 1] and [0, 1, 1] are subarrays, but [1, 1, 1] (skipping the 0) is not.

Since the answer can be very large, you need to return the result modulo 10^9 + 7.

Example breakdown:

• If we have a subarray [1, 0, 1, 1], it contains three 1s and one 0, so the count of 1s (3) is greater than the count of 0s (1). This subarray would be counted.
• If we have a subarray [0, 0, 1], it contains one 1 and two 0s, so the count of 1s (1) is not greater than the count of 0s (2). This subarray would not be counted.

The challenge is to efficiently count all such valid subarrays across all possible subarrays of the given array.

Intuition

The key insight is to transform this counting problem into something more manageable. Instead of directly counting 1s and 0s in every subarray, we can use a clever transformation: treat each 0 as -1 and keep 1 as 1. Now the problem becomes: count subarrays where the sum is greater than 0.

Why does this transformation help? With this change, a subarray has more 1s than 0s if and only if its sum is positive. For example, [1, 0, 1, 1] becomes [1, -1, 1, 1] with sum 2 > 0, indicating more 1s than 0s.

To count subarrays with positive sum, we use prefix sums. If prefix[i] is the sum of elements from index 0 to i, then the sum of subarray from index j+1 to i equals prefix[i] - prefix[j]. We want this difference to be positive, meaning prefix[i] > prefix[j].

So for each position i, we need to count how many previous positions j (where j < i) have prefix[j] < prefix[i]. This is essentially asking: "How many prefix sums before position i are smaller than the current prefix sum?"

This "count how many values are less than X" problem is perfectly suited for a Binary Indexed Tree (Fenwick Tree), which can efficiently:

1. Query the count of values in a range
2. Update counts as we process new elements

We process the array from left to right, maintaining the frequency of each prefix sum seen so far. For each new position, we query how many smaller prefix sums exist (these form valid subarrays ending at current position), then add the current prefix sum to our data structure.

The base offset in the code handles negative prefix sums by shifting all values to be positive, since Binary Indexed Trees typically work with positive indices.

Learn more about Segment Tree, Binary Search, Divide and Conquer and Merge Sort patterns.

Solution Approach

The solution uses Prefix Sum combined with a Binary Indexed Tree (Fenwick Tree) to efficiently count subarrays.

Step 1: Transform the Problem

• Convert each 0 to -1 in the calculation (not modifying the original array)
• This is done with: s += x or -1 where x or -1 returns 1 if x=1, and -1 if x=0
• Now we need to count subarrays with sum > 0

Step 2: Initialize Data Structures

• Create a Binary Indexed Tree with size n + base where base = n + 1
• The base offset ensures all indices are positive (handling negative prefix sums)
• Initialize the tree with prefix sum 0 having count 1 by calling tree.update(base, 1)
• This represents the empty prefix before any elements

Step 3: Process Each Element

• Maintain running prefix sum s and answer counter ans
• For each element in nums:
  1. Update prefix sum: s += x or -1
  2. Query how many previous prefix sums are less than current s:
     • Call tree.query(s - 1 + base) to get count of all prefix sums < s
     • These form valid subarrays ending at current position
  3. Add this count to ans (with modulo operation)
  4. Update the tree to include current prefix sum: tree.update(s + base, 1)

Binary Indexed Tree Operations:

• Update(x, v): Adds value v to position x
  • Uses bit manipulation x += x & -x to traverse parent nodes
  • Time complexity: O(log n)
• Query(x): Returns sum of values from position 1 to x
  • Uses bit manipulation x -= x & -x to traverse the tree
  • Time complexity: O(log n)

Why This Works:

• At position i, we've seen prefix sums from all positions 0 to i-1
• Querying for prefix sums less than s gives us all valid starting positions for subarrays ending at i
• Each such position j creates a subarray [j+1...i] with positive sum

Time Complexity: O(n log n) - we perform O(log n) operations for each of the n elements Space Complexity: O(n) - for the Binary Indexed Tree

Example Walkthrough

Let's walk through the solution with nums = [1, 0, 1, 0].

Step 1: Transform to +1/-1

• Original: [1, 0, 1, 0]
• Transform 0s to -1s: [1, -1, 1, -1]

Step 2: Initialize

• Create BIT with size n + base = 4 + 5 = 9 (base = n + 1 = 5)
• Initialize with prefix sum 0: tree.update(5, 1) (0 + base = 5)
• Set s = 0 (running prefix sum), ans = 0 (result counter)

Step 3: Process each element

Position 0: nums[0] = 1

• Update prefix sum: s = 0 + 1 = 1
• Query: How many prefix sums < 1?
  • tree.query(1 - 1 + 5) = tree.query(5) returns 1
  • This means 1 subarray ending here has positive sum: [1]
• Update answer: ans = 0 + 1 = 1
• Update tree: tree.update(1 + 5, 1) to record prefix sum 1

Position 1: nums[1] = 0 (transformed to -1)

• Update prefix sum: s = 1 + (-1) = 0
• Query: How many prefix sums < 0?
  • tree.query(0 - 1 + 5) = tree.query(4) returns 0
  • No subarrays ending here have positive sum
• Update answer: ans = 1 + 0 = 1
• Update tree: tree.update(0 + 5, 1) to record prefix sum 0

Position 2: nums[2] = 1

• Update prefix sum: s = 0 + 1 = 1
• Query: How many prefix sums < 1?
  • tree.query(1 - 1 + 5) = tree.query(5) returns 2
  • Two subarrays ending here have positive sum:
    • From after prefix 0 at start: [1, 0, 1] (sum = 1)
    • From after prefix 0 at position 1: [1] (sum = 1)
• Update answer: ans = 1 + 2 = 3
• Update tree: tree.update(1 + 5, 1) to record prefix sum 1

Position 3: nums[3] = 0 (transformed to -1)

• Update prefix sum: s = 1 + (-1) = 0
• Query: How many prefix sums < 0?
  • tree.query(0 - 1 + 5) = tree.query(4) returns 0
  • No subarrays ending here have positive sum
• Update answer: ans = 3 + 0 = 3
• Update tree: tree.update(0 + 5, 1) to record prefix sum 0

Final Answer: 3

The valid subarrays are:

1. [1] at position 0 (1 > 0)
2. [1] at position 2 (1 > 0)
3. [1, 0, 1] from positions 0-2 (2 ones > 1 zero)

Solution Implementation

Time and Space Complexity

Time Complexity: O(n × log n)

The main solution iterates through the array once with n iterations. For each element in the array:

• The query operation on the Binary Indexed Tree takes O(log n) time, as it traverses up the tree by repeatedly subtracting the least significant bit (x -= x & -x), which can happen at most log n times
• The update operation also takes O(log n) time, as it traverses up the tree by repeatedly adding the least significant bit (x += x & -x), which can happen at most log n times

Since we perform both query and update operations for each of the n elements, the overall time complexity is O(n × log n).

Space Complexity: O(n)

The space usage consists of:

• The Binary Indexed Tree's array c with size n + base + 1, where base = n + 1, resulting in approximately 2n + 2 elements, which is O(n)
• A few constant variables (base, mod, ans, s) which take O(1) space

Therefore, the total space complexity is O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Incorrect Problem Interpretation

The code provided actually solves for subarrays with more zeros than ones, but the problem statement asks for subarrays with more ones than zeros. This is a critical mismatch.

The Issue:

• The transformation prefix_sum += 1 if num else -1 makes zeros contribute -1 and ones contribute +1
• Querying for prefix sums less than current (fenwick_tree.query(prefix_sum - 1 + offset)) counts subarrays with negative sum
• Negative sum means more zeros than ones, not more ones than zeros

Solution: Change the transformation to make zeros contribute +1 and ones contribute -1, or change the query logic:

# Option 1: Flip the transformation
prefix_sum += -1 if num else 1  # Now 1 becomes -1, 0 becomes +1

# Option 2: Keep transformation but query differently
# Count prefix sums greater than current (requires range query)
# This is more complex with basic BIT

2. Offset Calculation Confusion

The offset is used to handle negative prefix sums, but it's easy to miscalculate the required range.

The Issue:

• With n elements, prefix sum can range from -n to +n
• Current offset of n + 1 may not be sufficient for all cases
• Incorrect offset can cause array index out of bounds or incorrect counting

Solution: Use a safer offset calculation:

# Maximum possible negative prefix sum is -n (all zeros)
# Maximum possible positive prefix sum is +n (all ones)
offset = n + 1  # This ensures all indices are positive

3. Forgetting Initial State

Not initializing the BIT with the empty prefix can lead to missing valid subarrays that start from index 0.

The Issue:

• Without fenwick_tree.update(offset, 1), subarrays starting from the beginning aren't counted
• For example, if the first element gives a valid subarray [0:1], it won't be counted

Solution: Always initialize with the empty prefix:

# This represents prefix sum of 0 before processing any elements
fenwick_tree.update(offset, 1)

4. Modulo Operation Timing

Applying modulo only at the end or forgetting it during intermediate steps can cause integer overflow in languages with fixed integer sizes.

The Issue:

• While Python handles arbitrary precision integers, forgetting modulo can still impact performance
• In other languages, this would cause overflow

Solution: Apply modulo after each addition:

result = (result + count_valid_subarrays) % MOD

Corrected Solution for "More Ones than Zeros":

def subarraysWithMoreOnesThanZeros(self, nums: List[int]) -> int:
    n = len(nums)
    offset = n + 1
    fenwick_tree = BinaryIndexedTree(2 * n + 1)  # Ensure sufficient size
  
    # Initialize with empty prefix
    fenwick_tree.update(offset, 1)
  
    MOD = 10**9 + 7
    result = 0
    prefix_sum = 0
  
    for num in nums:
        # Transform: 1 stays 1, 0 becomes -1
        prefix_sum += 1 if num else -1
      
        # Count all previous prefix sums that are LESS than current
        # For positive sum subarrays (more 1s than 0s)
        count_valid = fenwick_tree.query(prefix_sum - 1 + offset)
        result = (result + count_valid) % MOD
      
        # Record current prefix sum
        fenwick_tree.update(prefix_sum + offset, 1)
  
    return result