LeetCode 2031 - Count Subarrays With More Ones Than Zeros

This problem asks us to count all subarrays of a binary array nums that contain more 1s than 0s. In other words, for any contiguous slice of the array, if the number of 1s exceeds the number of 0s, it should be counted.

leetcodemediumarrayhash-tablebinary-searchdivide-and-conquerbinary-indexed-treesegment-treemerge-sortordered-set

LeetCode Problem 2031

Difficulty: 🟡 Medium
Topics: Array, Hash Table, Binary Search, Divide and Conquer, Binary Indexed Tree, Segment Tree, Merge Sort, Ordered Set

Solution

Problem Understanding

This problem asks us to count all subarrays of a binary array nums that contain more 1s than 0s. In other words, for any contiguous slice of the array, if the number of 1s exceeds the number of 0s, it should be counted. The input array only contains 0s and 1s, and the output must be returned modulo $10^9 + 7$ because the number of valid subarrays can be very large.

The constraints tell us that nums can have up to $10^5$ elements, so a brute-force approach that examines all possible subarrays (which would be $O(n^2)$) is likely too slow. Edge cases to consider include arrays with all 0s, all 1s, or a single element.

Approaches

The brute-force approach is straightforward: iterate over all possible subarrays, count the number of 1s and 0s in each, and increment a counter whenever 1s exceed 0s. This guarantees correctness because it directly checks each subarray, but it is prohibitively slow for large inputs due to $O(n^2)$ subarrays and $O(n)$ counting per subarray in a naive implementation, leading to $O(n^3)$ total time complexity.

The optimal approach relies on a key insight: we can transform the array by mapping 1 → +1 and 0 → -1. The problem then reduces to counting subarrays whose sum is positive. Using prefix sums, the sum of any subarray from index i to j is prefix[j+1] - prefix[i]. We can efficiently count how many previous prefix sums are smaller than the current prefix sum using a Fenwick Tree (Binary Indexed Tree) or a balanced ordered set.

The steps are: transform the array, compute prefix sums, and maintain a data structure that allows querying the number of prefix sums less than the current one. This reduces time complexity to roughly $O(n \log n)$.

Approach Time Complexity Space Complexity Notes
Brute Force O(n^2) O(1) Checks all subarrays directly, too slow for n=10^5
Optimal O(n log n) O(n) Transform to prefix sums and use BIT or ordered set to count positives efficiently

Algorithm Walkthrough

  1. Transform the binary array nums such that each 0 becomes -1 and each 1 remains 1. This makes the problem equivalent to counting subarrays with a positive sum.
  2. Initialize a prefix sum array, starting with 0. This helps compute sums of subarrays efficiently.
  3. Use a Fenwick Tree (Binary Indexed Tree) to maintain a frequency count of prefix sums seen so far.
  4. For each element in the transformed array, update the current prefix sum and query the Fenwick Tree to count how many previous prefix sums are less than the current prefix sum. This gives the number of subarrays ending at the current index that have a positive sum.
  5. Update the Fenwick Tree to include the current prefix sum.
  6. Accumulate these counts modulo $10^9 + 7$ and return the result.

Why it works: Every subarray sum prefix[j+1] - prefix[i] > 0 is counted by querying the number of prefix sums prefix[i] < prefix[j+1]. The Fenwick Tree allows us to maintain and query these counts efficiently in logarithmic time.

Python Solution

from typing import List

class FenwickTree:
    def __init__(self, size: int):
        self.tree = [0] * (size + 2)
        self.size = size + 2
    
    def update(self, index: int, delta: int):
        while index < self.size:
            self.tree[index] += delta
            index += index & -index
    
    def query(self, index: int) -> int:
        res = 0
        while index > 0:
            res += self.tree[index]
            index -= index & -index
        return res

class Solution:
    def subarraysWithMoreOnesThanZeroes(self, nums: List[int]) -> int:
        MOD = 10**9 + 7
        n = len(nums)
        arr = [1 if num == 1 else -1 for num in nums]
        prefix = [0]
        for num in arr:
            prefix.append(prefix[-1] + num)
        
        # Coordinate compression
        sorted_prefix = sorted(set(prefix))
        rank = {val: i+1 for i, val in enumerate(sorted_prefix)}
        
        ft = FenwickTree(len(rank))
        count = 0
        for p in prefix:
            r = rank[p]
            # Count of prefix sums < current prefix sum
            count += ft.query(r - 1)
            count %= MOD
            ft.update(r, 1)
        return count

Implementation Explanation: We first transform nums into +1/-1 values, then compute prefix sums. Coordinate compression maps prefix sums to smaller indices suitable for the Fenwick Tree. For each prefix sum, querying the number of previous prefix sums smaller than it gives the number of positive subarrays ending at that point. The Fenwick Tree is updated with the current prefix sum after counting.

Go Solution

package main

func lowbit(x int) int { return x & -x }

type FenwickTree struct {
    tree []int
    size int
}

func NewFenwickTree(size int) *FenwickTree {
    return &FenwickTree{tree: make([]int, size+2), size: size + 2}
}

func (ft *FenwickTree) Update(index, delta int) {
    for index < ft.size {
        ft.tree[index] += delta
        index += lowbit(index)
    }
}

func (ft *FenwickTree) Query(index int) int {
    res := 0
    for index > 0 {
        res += ft.tree[index]
        index -= lowbit(index)
    }
    return res
}

func subarraysWithMoreOnesThanZeroes(nums []int) int {
    MOD := 1000000007
    n := len(nums)
    arr := make([]int, n)
    for i, num := range nums {
        if num == 1 {
            arr[i] = 1
        } else {
            arr[i] = -1
        }
    }

    prefix := make([]int, n+1)
    for i := 0; i < n; i++ {
        prefix[i+1] = prefix[i] + arr[i]
    }

    // Coordinate compression
    m := make(map[int]struct{})
    for _, p := range prefix {
        m[p] = struct{}{}
    }
    sorted := make([]int, 0, len(m))
    for key := range m {
        sorted = append(sorted, key)
    }
    sort.Ints(sorted)
    rank := make(map[int]int)
    for i, val := range sorted {
        rank[val] = i + 1
    }

    ft := NewFenwickTree(len(rank))
    count := 0
    for _, p := range prefix {
        r := rank[p]
        count = (count + ft.Query(r-1)) % MOD
        ft.Update(r, 1)
    }
    return count
}

Go-specific Notes: Go requires explicit map creation for coordinate compression. Slices are used instead of arrays for dynamic size handling. Integer modulo is applied to avoid overflow.

Worked Examples

Example 1: nums = [0,1,1,0,1]

Index Num Transformed Prefix Count query Fenwick update
0 0 -1 0 0 add 1 at rank(0)=2
1 1 1 0 0 add 1 at rank(0)=2
2 1 1 1 2 add 1 at rank(1)=4
3 0 -1 0 1 add 1 at rank(0)=2
4 1 1 1 3 add 1 at rank(1)=4
5 - - 2 3 add 1 at rank(2)=5

Total count = 9

Example 2: nums = [0]

Index Num Transformed Prefix Count query Fenwick update
0 0 -1 0 0 add 1

Total count = 0

Example 3: nums = [1]

Index Num Transformed Prefix Count query Fenwick update
0 1 1 1 1 add 1

Total count = 1

Complexity Analysis

Measure Complexity Explanation
Time O(n log n) Transform array is O(n), prefix sum computation O(n), coordinate compression O(n log n), Fenwick operations O(log n) per prefix sum for n elements