2031. Count Subarrays With More Ones Than Zeros

MediumBinary Indexed TreeSegment TreeArrayBinary SearchDivide and ConquerOrdered SetMerge Sort

Problem Description

In this problem, we are given a binary array nums containing only integers 0 and 1. The task is to determine how many contiguous subarrays exist within the given array where the number of 1s is greater than the number of 0s. Due to potentially large output, the result should be returned modulo 10^9 + 7.

A subarray is a sequence of consecutive elements from the array, which implies that each element of the array nums can be the start or endpoint of multiple subarrays. If we were to count each possible subarray naively, checking if 1s are more than 0s in all of them, this would lead to a very high time complexity exceeding the limits of most computational tasks.

The intuition behind solving this problem efficiently is to use a data structure that can help us keep track of the count of 1s and 0s as we move through the array. The prefix sum technique can be applied here to help with this. Since 1s increase the sum and 0s decrease it, we can convert 0s to -1s and use prefix sums to find the number of 1s and 0s encountered up to every point in the array.

Intuition

The solution employs a Binary Indexed Tree (BIT), also known as a Fenwick Tree, which is an efficient data structure for maintaining and querying prefix sums of a list of numbers.

Here's the thought process leading to the solution:

Prefix Sum Transformation: Transform the array into a prefix sum array where every element s[i] is the sum of values from nums[0] to nums[i], replacing each 0 in nums with -1. This allows us to find the number of 1s and 0s seen so far as we iterate along the array.
Understanding Subarrays with More 1s than 0s: A subarray has more 1s than 0s if the prefix sum at its end is greater than the prefix sum at its start. This corresponds to a positive change in the sum, indicating more 1s have been added than 0s (or -1s when transformed).
Counting With Binary Indexed Tree: Use the BIT to keep count of prefix sums seen so far. As we iterate over the transformed array's prefix sums, we query the BIT for the number of times a prefix sum less than the current one has occurred. This corresponds to counting the number of previously encountered subarrays where more 1s than 0s would ensure a positive sum difference, indicating more 1s.
Handling Negative Sums and Offsetting: Since prefix sums can become negative (more 0s seen at some point), and the BIT cannot handle negative indices, an offset is added to all sums to ensure they are non-negative.
Updating and Querying the BIT: Every time we encounter a new sum, we update the BIT with this sum, incrementing the count at that index. We also query the BIT for the total counts of all sums less than (but not including) the current sum to get the number of subarrays ending at the current index with more 1s than 0s.
Modulo Operation: Since the answer can be very large, every time a new count is added to the result, it is taken modulo 10^9 + 7 to keep it within the required range.

By using the BIT to efficiently keep track of prefix sums and their frequencies, the solution can iterate through the array just once and count the number of desired subarrays with the criteria given.

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

Solution Approach

The implementation employs a BinaryIndexedTree class, which encapsulates the functionality of a Fenwick Tree to perform efficient updates and queries on the prefix sums.

Binary Indexed Tree (Fenwick Tree)

A Fenwick Tree supports two operations efficiently: updating the value of an element (in this context, it's the frequency count of a particular prefix sum) and querying the prefix sum of a sequence up to a certain point (in this context, it's to find out how many times a prefix sum less than the current one has occurred).

Class Implementation Details:

The BinaryIndexedTree class has a constructor which initializes an array self.c of size n + 1, filled with zeroes. Here n is the size of the input array plus an offset to handle negative sums.
The lowbit function is a static method that takes a number x and returns the largest power of 2 that divides x (bitwise AND of x with its two's complement).
The update function increments the count of a particular prefix sum in the BIT. It takes an index x and a value delta. This function iterates through the tree updating all relevant nodes affected by the insertion of this new sum, augmenting them by delta.
The query function calculates the cumulative frequency counts up to index x. It essentially computes the number of elements that have a prefix sum which is less than the current prefix sum.

Solution Class - `subarraysWithMoreZerosThanOnes` Method:

The core function subarraysWithMoreZerosThanOnes starts by initializing the prefix sum array s with a zero to handle the empty subarray case.
It iterates over the nums array to build the prefix sum list by adding 1 or -1 depending on whether the current value is a 1 or 0, respectively.
The BinaryIndexedTree is initialized to handle the prefix sums.
MOD is defined as 10^9 + 7 to ensure the results are within the range.
The answer ans is initially set to 0.
For every prefix sum v in s, it queries the tree for the number of prefix sums less than v (which corresponds to v - 1) and adds this number to ans.
After each query, it updates the tree with the current prefix sum v to increment its count.
It uses modulo operation for each sum addition to ans to prevent overflow and meet the problem's requirement of returning the result modulo 10^9 + 7.

This use of a Binary Indexed Tree enables the solution to process each element of the array in logarithmic time with respect to the size of the input array, which is significantly faster than a brute force solution that would have quadratic time complexity.

Ready to land your dream job?

Unlock your dream job with a 2-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through a small example to illustrate the solution approach using the array nums = [1,0,1,1,0].

Step-by-Step Process

Prefix Sum Transformation: Transform nums into a prefix sum array s. First, we initialize s with a starting value of 0. Then we iterate over nums, adding 1 for 1s and -1 for 0s.
- Start with s = [0].
- After the first element (1), s = [0, 1].
- Next element (0), s = [0, 1, 0]; here, we added -1.
- Next element (1), s = [0, 1, 0, 1].
- Next element (1), s = [0, 1, 0, 1, 2].
- Last element (0), s = [0, 1, 0, 1, 2, 1].
Using Binary Indexed Tree: Initialize a Binary Indexed Tree (BIT) with enough space to handle all possible prefix sums, including the offset.
Offsetting: Since BIT cannot handle negative indices and our prefix sums may include negatives (though not in this specific example), we would typically add an offset. However, in this example, there are no negative prefix sums, so we can skip this step.
Iterating and Counting: Iterate over the transformed prefix sum array s while querying and updating the BIT.
- Start with ans = 0.
- At the first prefix sum 1:
  - Query the BIT for counts of sums less than 1, which is 0. Thus, we add 0 to ans.
  - Update the BIT with the sum 1.
- At the prefix sum 0:
  - Query the BIT for counts of sums less than 0, which is 0 (no sums are less than zero). We add 0 to ans.
  - Update the BIT with the sum 0.
- At the prefix sum 1 again:
  - Query the BIT for counts of sums less than 1, which is now 1 (due to the previous update). We add 1 to ans, making ans = 1.
  - Update the BIT with the sum 1 again.
- At the prefix sum 2:
  - Query the BIT for counts of sums less than 2, which is 2 (from two previous 1s). We add 2 to ans, so now ans = 3.
  - Update the BIT with the sum 2.
- At the last prefix sum 1 again:
  - Query the BIT for counts of sums less than 1, now 2 from the previous occurrences. We add 2 to ans and get ans = 5.
  - Update the BIT with 1.
Modulo Operation: In this example, the answer does not exceed the limit, so we use ans as it is. If it were larger, we would use ans % MOD where MOD = 10^9 + 7.

Our final ans is 5, representing the number of contiguous subarrays in nums where the number of 1s is greater than the number of 0s. The subarrays that satisfy this condition are:

[1]
[1,0,1]
[1,0,1,1]
[0,1,1]
[1,1]

Hence, the approach leverages the efficient update and query capabilities of the BIT to compute the number of subarrays with more 1s than 0s.

Solution Implementation

1class BinaryIndexedTree:
2    def __init__(self, n):
3        # Adjust the size based on the problem offset
4        self.size = n + int(1e5 + 1)
5        # Initialize the binary indexed tree with zeros
6        self.tree = [0] * (self.size + 1)
7
8    def update(self, index, delta):
9        # Adjust the index for the problem offset
10        index += int(1e5 + 1)
11        # Update the tree by adding 'delta' to each node along the path
12        while index <= self.size:
13            self.tree[index] += delta
14            # Move to the next index to update
15            # lowbit function is inlined for performance
16            index += index & -index
17
18    def query(self, index):
19        # Adjust the index for the problem offset
20        index += int(1e5 + 1)
21        result = 0
22        # Sum the values in the tree along the path to the root node
23        while index > 0:
24            result += self.tree[index]
25            # Move to the parent index to continue the sum
26            # lowbit function is inlined for performance
27            index -= index & -index
28        return result
29
30
31class Solution:
32    def subarraysWithMoreZerosThanOnes(self, nums):
33        # Calculate the length of the given 'nums' array
34        length = len(nums)
35        prefix_sums = [0]
36        # Compute prefix sums, decrementing for zeros and incrementing for ones
37        for value in nums:
38            prefix_sums.append(prefix_sums[-1] + (value or -1))
39      
40        # Instantiate a `BinaryIndexedTree` with length of our prefix sums
41        tree = BinaryIndexedTree(length + 1)
42        # Define the modulo for result to keep it within bounds
43        MOD = int(1e9 + 7)
44        # Initialize answer to 0
45        answer = 0
46      
47        # Iterate over the prefix sums
48        for value in prefix_sums:
49            # Query the tree for the range up to value - 1
50            # and update the answer considering MOD
51            answer = (answer + tree.query(value - 1)) % MOD
52            # Update the tree with the current prefix sum value
53            tree.update(value, 1)
54      
55        # Return the final answer
56        return answer
57

1class BinaryIndexedTree {
2    private int size;
3    private int[] treeArray;
4
5    // Constructor to initialize the Binary Indexed Tree with a given size
6    // The size is increased by an offset to handle negative indices
7    public BinaryIndexedTree(int size) {
8        int offset = (int) 1e5 + 1;
9        this.size = size + offset;
10        this.treeArray = new int[this.size + 1];
11    }
12
13    // Updates the Binary Indexed Tree at a specific position 'x' with value 'delta'
14    public void update(int x, int delta) {
15        int offset = (int) 1e5 + 1;
16        x += offset;
17        while (x <= size) {
18            treeArray[x] += delta;
19            x += lowbit(x); // Moves to the next index to be updated
20        }
21    }
22
23    // Queries the sum of values within the range [1, x] in the Binary Indexed Tree
24    public int query(int x) {
25        int offset = (int) 1e5 + 1;
26        x += offset;
27        int sum = 0;
28        while (x > 0) {
29            sum += treeArray[x];
30            x -= lowbit(x); // Moves down the tree to sum up the values
31        }
32        return sum;
33    }
34
35    // Returns the least significant bit of 'x'
36    // The offset here should be ignored as it's a static utility function
37    public static int lowbit(int x) {
38        return x & -x;
39    }
40}
41
42class Solution {
43    private static final int MOD = (int) 1e9 + 7; // Modulus for avoiding overflow
44
45    // Counts the number of subarrays with more zeros than ones in a binary array
46    public int subarraysWithMoreZerosThanOnes(int[] nums) {
47        int n = nums.length; // Length of the input array
48        int[] prefixSums = new int[n + 1]; // Array for storing prefix sums
49
50        // Calculate the prefix sums array
51        for (int i = 0; i < n; ++i) {
52            prefixSums[i + 1] = prefixSums[i] + (nums[i] == 1 ? 1 : -1);
53        }
54
55        BinaryIndexedTree tree = new BinaryIndexedTree(n + 1);
56        int answer = 0; // Initialize result
57      
58        // Iterate over the prefix sums array while updating and querying the tree
59        for (int value : prefixSums) {
60            // Query the number of indices with prefix sums less than the current
61            // and update the answer
62            answer = (answer + tree.query(value - 1)) % MOD;
63            // Update the tree's count for the current prefix sum
64            tree.update(value, 1);
65        }
66        return answer;
67    }
68}
69

1class BinaryIndexedTree {
2public:
3    int size;                // Size of the original array
4    vector<int> bit;         // Binary Indexed Tree storage
5
6    // Constructor with an argument _n representing the size of the array
7    // Adding a large number (1e5 + 1) to account for negative indexes
8    // because prefix sums can be negative, but BIT is 1-indexed and must be positive
9    BinaryIndexedTree(int _n)
10        : size(_n + static_cast<int>(1e5) + 1)
11        , bit(_n + static_cast<int>(1e5) + 2) {}
12
13    // Updates the BIT at a given index 'x' by a certain value 'delta'
14    void update(int x, int delta) {
15        x += static_cast<int>(1e5) + 1; // Adjust index to be positive
16        while (x <= size) {
17            bit[x] += delta; // Update the BIT
18            x += lowbit(x);  // Move to the next index to be updated
19        }
20    }
21
22    // Query the BIT up to a certain index 'x'
23    int query(int x) {
24        x += static_cast<int>(1e5) + 1; // Adjust index to be positive
25        int sum = 0;
26        while (x > 0) {
27            sum += bit[x];  // Add current element to the sum
28            x -= lowbit(x); // Move to the previous index in BIT
29        }
30        return sum;
31    }
32
33    // Helper function to calculate the lowest bit of a given number
34    // which represents the power of two that partitions the segments stored in BIT
35    int lowbit(int x) {
36        return x & -x;
37    }
38};
39
40class Solution {
41public:
42    // Function to count the number of subarrays with more zeros than ones
43    int subarraysWithMoreZerosThanOnes(vector<int>& nums) {
44        int n = nums.size();
45        // Prefix sums array with an extra element for the starting 0
46        vector<int> prefixSums(n + 1);
47
48        // Calculate the prefix sums array wherein 
49        // 1s are counted as 1 and 0s as -1 to later find subarrays with negative sums
50        for (int i = 0; i < n; ++i) {
51            prefixSums[i + 1] = prefixSums[i] + (nums[i] == 1 ? 1 : -1);
52        }
53
54        // Allocate a BIT on heap to store cumulative counts
55        BinaryIndexedTree* tree = new BinaryIndexedTree(n + 1);
56        int count = 0; // This will store the final answer
57        const int MOD = static_cast<int>(1e9) + 7;
58
59        for (int value : prefixSums) {
60            // Query the count of subarrays ending before 'value' (exclusive)
61            count = (count + tree->query(value - 1)) % MOD;
62            // Update the BIT with the current prefix sum value
63            tree->update(value, 1);
64        }
65
66        // Since 'tree' was allocated using 'new', it should be deleted to prevent memory leak
67        delete tree;
68
69        return count;
70    }
71};
72

1// The offset to make sure all indices are positive, since Binary Indexed Tree is 1-indexed
2const OFFSET = 100001;
3
4// The size of the nums array to be used for initializing the BIT array
5let size: number;
6
7// The Binary Indexed Tree storage array
8let bit: number[];
9
10// Initialize the BIT array to a certain size
11function initBIT(n: number): void {
12    size = n + OFFSET;
13    bit = new Array(size + 1).fill(0);
14}
15
16// Get the lowest bit of a given number
17function lowbit(x: number): number {
18    return x & -x;
19}
20
21// Update the BIT at a given index 'x' with 'delta'
22function update(x: number, delta: number): void {
23    x += OFFSET; // Adjust index to be positive
24    while (x <= size) {
25        bit[x] += delta; // Update the BIT
26        x += lowbit(x); // Move to the next index to be updated
27    }
28}
29
30// Query the BIT up to a certain index 'x'
31function query(x: number): number {
32    x += OFFSET; // Adjust index to be positive
33    let sum = 0;
34    while (x > 0) {
35        sum += bit[x]; // Add current element to the sum
36        x -= lowbit(x); // Move to the previous index in BIT
37    }
38    return sum;
39}
40
41// Count the number of subarrays with more zeros than ones
42function subarraysWithMoreZerosThanOnes(nums: number[]): number {
43    const n: number = nums.length;
44    const prefixSums: number[] = new Array(n + 1);
45    prefixSums.fill(0);
46  
47    // Calculate the prefix sums where 1s are 1 and 0s are -1
48    for (let i = 0; i < n; ++i) {
49        prefixSums[i + 1] = prefixSums[i] + (nums[i] === 1 ? 1 : -1);
50    }
51  
52    // Initialize the BIT with extra space for negative indices
53    initBIT(n + 1);
54  
55    let count: number = 0; // To store the final answer
56    const MOD: number = 1_000_000_007;
57
58    for (let value of prefixSums) {
59        // Query the count of subarrays ending before 'value' (exclusive)
60        count = (count + query(value - 1)) % MOD;
61        // Update the BIT with the current prefix sum value
62        update(value, 1);
63    }
64
65    return count;
66}
67

Time and Space Complexity

Time Complexity

The time complexity of the code is determined by several key parts: the preprocessing step where the prefix sum array s is constructed, the binary indexed tree operations (update and query), and the final loop where we use the tree to find the number of subarrays with more zeros than ones.

Constructing the prefix sum array s requires a single pass through the array nums, which takes O(n) time.
The update and query operations of the binary indexed tree have a time complexity of O(log n) each since each operation involves traversing the tree up to a depth proportional to the logarithm of the number of elements.
The final loop runs n + 1 times, and each iteration performs an update and a query operation on the binary indexed tree. This makes the complexity of the loop O(n log n).

Therefore, the total time complexity of the code is O(n) + O(n log n) = O(n log n).

Space Complexity

The space complexity is determined by the size of the data structures used:

The prefix sum array s has n + 1 elements, so it takes O(n) space.
The binary indexed tree object holds n + 1 elements, considering the offset for handling negative values in the array nums. Since we add int(1e5 + 1) to handle the range of values, the actual size of the array c within the tree is n + int(1e5 + 1) + 1, which is O(n) as n tends to infinity.

Hence, the space complexity of the code is O(n).

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

Discover Your Strengths and Weaknesses: Take Our 2-Minute Quiz to Tailor Your Study Plan:

Question 1 out of 10

What is the best way of checking if an element exists in a sorted array once in terms of time complexity? Select the best that applies.

Linear Search

Binary Search

Quick Select

Hash Set

I don’t know