Welcome to Subscribe On Youtube
2031. Count Subarrays With More Ones Than Zeros
Description
You are given a binary array nums containing only the integers 0 and 1. Return the number of subarrays in nums that have more 1's than 0's. Since the answer may be very large, return it modulo 109 + 7.
A subarray is a contiguous sequence of elements within an array.
Example 1:
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:
Input: nums = [0] Output: 0 Explanation: No subarrays have more ones than zeros.
Example 3:
Input: nums = [1] Output: 1 Explanation: The subarrays of size 1 that have more ones than zeros are: [1]
Constraints:
1 <= nums.length <= 1050 <= nums[i] <= 1
Solutions
Solution 1: Prefix Sum + Binary Indexed Tree
The problem requires us to count the number of subarrays where the count of $1$ is greater than the count of $0$. If we treat $0$ in the array as $-1$, then the problem becomes counting the number of subarrays where the sum of elements is greater than $0$.
To calculate the sum of elements in a subarray, we can use the prefix sum. To count the number of subarrays where the sum of elements is greater than $0$, we can use a binary indexed tree to maintain the occurrence count of each prefix sum. Initially, the occurrence count of the prefix sum $0$ is $1$.
Next, we traverse the array $nums$, use variable $s$ to record the current prefix sum, and use variable $ans$ to record the answer. For each position $i$, we update the prefix sum $s$, then query the occurrence count of the prefix sum in the range $[0, s)$ in the binary indexed tree, add it to $ans$, and then update the occurrence count of $s$ in the binary indexed tree.
Finally, return $ans$.
The time complexity is $O(n \times \log n)$, and the space complexity is $O(n)$. Where $n$ is the length of the array $nums$.
Solution 2
This implementation uses binary search. It traverses the relevant values and updates its state as each value is processed. Ordering the data first makes the subsequent comparisons deterministic. After all required states have been considered, the maintained result is returned.
-
class BinaryIndexedTree { private int n; private int[] c; public BinaryIndexedTree(int n) { this.n = n; c = new int[n + 1]; } public void update(int x, int v) { for (; x <= n; x += x & -x) { c[x] += v; } } public int query(int x) { int s = 0; for (; x > 0; x -= x & -x) { s += c[x]; } return s; } } class Solution { public int subarraysWithMoreZerosThanOnes(int[] nums) { int n = nums.length; int base = n + 1; BinaryIndexedTree tree = new BinaryIndexedTree(n + base); tree.update(base, 1); final int mod = (int) 1e9 + 7; int ans = 0, s = 0; for (int x : nums) { s += x == 0 ? -1 : 1; ans += tree.query(s - 1 + base); ans %= mod; tree.update(s + base, 1); } return ans; } } -
class BinaryIndexedTree { private: int n; vector<int> c; public: BinaryIndexedTree(int n) : n(n) , c(n + 1, 0) {} void update(int x, int v) { for (; x <= n; x += x & -x) { c[x] += v; } } int query(int x) { int s = 0; for (; x > 0; x -= x & -x) { s += c[x]; } return s; } }; class Solution { public: int subarraysWithMoreZerosThanOnes(vector<int>& nums) { int n = nums.size(); int base = n + 1; BinaryIndexedTree tree(n + base); tree.update(base, 1); const int mod = 1e9 + 7; int ans = 0, s = 0; for (int x : nums) { s += (x == 0) ? -1 : 1; ans += tree.query(s - 1 + base); ans %= mod; tree.update(s + base, 1); } return ans; } }; -
class BinaryIndexedTree: __slots__ = ["n", "c"] def __init__(self, n: int): self.n = n self.c = [0] * (n + 1) def update(self, x: int, v: int): while x <= self.n: self.c[x] += v x += x & -x def query(self, x: int) -> int: s = 0 while x: s += self.c[x] x -= x & -x return s class Solution: def subarraysWithMoreZerosThanOnes(self, nums: List[int]) -> int: n = len(nums) base = n + 1 tree = BinaryIndexedTree(n + base) tree.update(base, 1) mod = 10**9 + 7 ans = s = 0 for x in nums: s += x or -1 ans += tree.query(s - 1 + base) ans %= mod tree.update(s + base, 1) return ans # Solution 2 class Solution: def subarraysWithMoreZerosThanOnes(self, nums: List[int]) -> int: sl = SortedList([0]) mod = 10**9 + 7 ans = s = 0 for x in nums: s += x or -1 ans += sl.bisect_left(s) ans %= mod sl.add(s) return ans -
type BinaryIndexedTree struct { n int c []int } func newBinaryIndexedTree(n int) *BinaryIndexedTree { return &BinaryIndexedTree{n: n, c: make([]int, n+1)} } func (bit *BinaryIndexedTree) update(x, v int) { for ; x <= bit.n; x += x & -x { bit.c[x] += v } } func (bit *BinaryIndexedTree) query(x int) (s int) { for ; x > 0; x -= x & -x { s += bit.c[x] } return } func subarraysWithMoreZerosThanOnes(nums []int) (ans int) { n := len(nums) base := n + 1 tree := newBinaryIndexedTree(n + base) tree.update(base, 1) const mod = int(1e9) + 7 s := 0 for _, x := range nums { if x == 0 { s-- } else { s++ } ans += tree.query(s - 1 + base) ans %= mod tree.update(s+base, 1) } return } -
class BinaryIndexedTree { private n: number; private c: number[]; constructor(n: number) { this.n = n; this.c = Array(n + 1).fill(0); } update(x: number, v: number): void { for (; x <= this.n; x += x & -x) { this.c[x] += v; } } query(x: number): number { let s = 0; for (; x > 0; x -= x & -x) { s += this.c[x]; } return s; } } function subarraysWithMoreZerosThanOnes(nums: number[]): number { const n: number = nums.length; const base: number = n + 1; const tree: BinaryIndexedTree = new BinaryIndexedTree(n + base); tree.update(base, 1); const mod: number = 1e9 + 7; let ans: number = 0; let s: number = 0; for (const x of nums) { s += x || -1; ans += tree.query(s - 1 + base); ans %= mod; tree.update(s + base, 1); } return ans; }