Problem

You are given a binary array nums containing only the integers 0 and 1. Return _the number ofsubarrays in nums that have more _1’s than0 _’ s. Since the answer may be very large, return it modulo _109 + 7.

A subarray is a contiguous sequence of elements within an array.

Examples

Example 1:

1
2
3
4
5
6
7
8

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:

1
2
3
4

Input: nums = [0]
Output: 0
Explanation:
No subarrays have more ones than zeros.

Example 3:

1
2
3
4

Input: nums = [1]
Output: 1
Explanation:
The subarrays of size 1 that have more ones than zeros are: [1]

Constraints:

  • 1 <= nums.length <= 10^5
  • 0 <= nums[i] <= 1

Solution

Method 1 – Prefix Sum and Binary Indexed Tree (Fenwick Tree)

Intuition

We can map 1 to +1 and 0 to -1, and use prefix sums to track the difference between ones and zeros. For each prefix sum, we want to count how many previous prefix sums are less than the current, which can be done efficiently with a Binary Indexed Tree (Fenwick Tree) after coordinate compression.

Approach

  1. Map 1 to +1 and 0 to -1, and compute prefix sums (starting from 0).
  2. Coordinate compress all possible prefix sums.
  3. Use a Fenwick Tree to count, for each prefix sum, how many previous prefix sums are less than the current.
  4. For each prefix sum, add the count to the answer and update the Fenwick Tree.
  5. Return the answer modulo 10^9+7.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30

class Solution {
    public int countSubarrays(int[] nums) {
        int n = nums.length, MOD = 1_000_000_7;
        int[] psum = new int[n+1];
        for (int i = 0; i < n; i++) psum[i+1] = psum[i] + (nums[i] == 1 ? 1 : -1);
        int[] sorted = psum.clone();
        Arrays.sort(sorted);
        Map<Integer, Integer> idx = new HashMap<>();
        int id = 0;
        for (int x : sorted) if (!idx.containsKey(x)) idx.put(x, id++);
        FenwickTree bit = new FenwickTree(id);
        int ans = 0;
        for (int x : psum) {
            int i = idx.get(x);
            ans = (ans + bit.query(i-1)) % MOD;
            bit.update(i, 1);
        }
        return ans;
    }
    class FenwickTree {
        int[] tree;
        FenwickTree(int n) { tree = new int[n+2]; }
        void update(int i, int v) { for (i++; i < tree.length; i += i & -i) tree[i] += v; }
        int query(int i) {
            int res = 0;
            for (i++; i > 0; i -= i & -i) res += tree[i];
            return res;
        }
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31

class Solution:
    def countSubarrays(self, nums: list[int]) -> int:
        MOD = 10**9 + 7
        n = len(nums)
        psum = [0]
        for x in nums:
            psum.append(psum[-1] + (1 if x == 1 else -1))
        sorted_psum = sorted(set(psum))
        idx = {v: i for i, v in enumerate(sorted_psum)}
        class BIT:
            def __init__(self, n):
                self.t = [0]*(n+2)
            def update(self, i, v):
                i += 1
                while i < len(self.t):
                    self.t[i] += v
                    i += i & -i
            def query(self, i):
                i += 1
                res = 0
                while i > 0:
                    res += self.t[i]
                    i -= i & -i
                return res
        bit = BIT(len(idx))
        ans = 0
        for x in psum:
            i = idx[x]
            ans = (ans + bit.query(i-1)) % MOD
            bit.update(i, 1)
        return ans

Complexity

  • ⏰ Time complexity: O(n log n), due to sorting and Fenwick Tree operations.
  • 🧺 Space complexity: O(n), for prefix sums and Fenwick Tree.