From c74f538b6bf8be24fa424fb326f468772f843373 Mon Sep 17 00:00:00 2001 From: yanglbme Date: Wed, 20 Dec 2023 09:51:10 +0800 Subject: [PATCH 1/2] feat: add solutions to lc problem: No.2031 No.2031.Count Subarrays With More Ones Than Zeros --- .../README.md | 239 +++++++++--------- .../README_EN.md | 238 +++++++++-------- .../Solution.cpp | 42 ++- .../Solution.go | 57 ++--- .../Solution.java | 39 +-- .../Solution.py | 43 ++-- .../Solution.ts | 40 +++ 7 files changed, 368 insertions(+), 330 deletions(-) create mode 100644 solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.ts diff --git a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README.md b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README.md index f5c73daf9ef58..7fb63bf340ea0 100644 --- a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README.md +++ b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README.md @@ -53,14 +53,17 @@ -**方法一：树状数组** +**方法一：前缀和 + 树状数组** -树状数组，也称作“二叉索引树”（Binary Indexed Tree）或 Fenwick 树。 它可以高效地实现如下两个操作： +题目需要我们统计所有子数组中 $1$ 的数量大于 $0$ 的数量的子数组的个数，如果我们将数组中的元素 $0$ 看作 $-1$，那么题目就变成了统计所有子数组中元素和大于 $0$ 的子数组的个数。 -1. **单点更新** `update(x, delta)`： 把序列 x 位置的数加上一个值 delta； -1. **前缀和查询** `query(x)`：查询序列 `[1,...x]` 区间的区间和，即位置 x 的前缀和。 +求子数组的元素和，可以使用前缀和来实现。为了统计所有子数组中元素和大于 $0$ 的子数组的个数，我们可以用树状数组维护每个前缀和出现的次数。初始时前缀和为 $0$ 的次数为 $1$。 -这两个操作的时间复杂度均为 $O(\log n)$。 +接下来，我们遍历数组 $nums$，用变量 $s$ 记录当前的前缀和，用变量 $ans$ 记录答案。对于每个位置 $i$，更新前缀和 $s$，然后我们在树状数组中查询 $[0, s)$ 范围内的前缀和出现的次数，将其加到 $ans$ 中，然后在树状数组中更新 $s$ 出现的次数。 + +最后返回 $ans$ 即可。 + +时间复杂度 $O(n \times \log n)$，空间复杂度 $O(n)$。其中 $n$ 是数组 $nums$ 的长度。 @@ -70,43 +73,38 @@ ```python class BinaryIndexedTree: - def __init__(self, n): - n += int(1e5 + 1) + __slots__ = ["n", "c"] + + def __init__(self, n: int): self.n = n self.c = [0] * (n + 1) - @staticmethod - def lowbit(x): - x += int(1e5 + 1) - return x & -x - - def update(self, x, delta): - x += int(1e5 + 1) + def update(self, x: int, v: int): while x <= self.n: - self.c[x] += delta - x += BinaryIndexedTree.lowbit(x) + self.c[x] += v + x += x & -x - def query(self, x): - x += int(1e5 + 1) + def query(self, x: int) -> int: s = 0 - while x > 0: + while x: s += self.c[x] - x -= BinaryIndexedTree.lowbit(x) + x -= x & -x return s class Solution: def subarraysWithMoreZerosThanOnes(self, nums: List[int]) -> int: n = len(nums) - s = [0] - for v in nums: - s.append(s[-1] + (v or -1)) - tree = BinaryIndexedTree(n + 1) - MOD = int(1e9 + 7) - ans = 0 - for v in s: - ans = (ans + tree.query(v - 1)) % MOD - tree.update(v, 1) + base = n + 1 + tree = BinaryIndexedTree(n + base) + tree.update(base, 1) + mod = 10**9 + 7 + ans = s = 0 + for x in nums: + s += 1 if x else -1 + ans += tree.query(s - 1 + base) + ans %= mod + tree.update(s + base, 1) return ans ``` @@ -120,49 +118,38 @@ class BinaryIndexedTree { private int[] c; public BinaryIndexedTree(int n) { - n += (int) 1e5 + 1; this.n = n; c = new int[n + 1]; } - public void update(int x, int delta) { - x += (int) 1e5 + 1; - while (x <= n) { - c[x] += delta; - x += lowbit(x); + public void update(int x, int v) { + for (; x <= n; x += x & -x) { + c[x] += v; } } public int query(int x) { - x += (int) 1e5 + 1; int s = 0; - while (x > 0) { + for (; x > 0; x -= x & -x) { s += c[x]; - x -= lowbit(x); } return s; } - - public static int lowbit(int x) { - x += (int) 1e5 + 1; - return x & -x; - } } class Solution { - private static final int MOD = (int) 1e9 + 7; - public int subarraysWithMoreZerosThanOnes(int[] nums) { int n = nums.length; - int[] s = new int[n + 1]; - for (int i = 0; i < n; ++i) { - s[i + 1] = s[i] + (nums[i] == 1 ? 1 : -1); - } - BinaryIndexedTree tree = new BinaryIndexedTree(n + 1); - int ans = 0; - for (int v : s) { - ans = (ans + tree.query(v - 1)) % MOD; - tree.update(v, 1); + 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; } @@ -173,50 +160,42 @@ class Solution { ```cpp class BinaryIndexedTree { -public: +private: int n; vector c; - BinaryIndexedTree(int _n) - : n(_n + 1e5 + 1) - , c(_n + 1 + 1e5 + 1) {} +public: + BinaryIndexedTree(int n) : n(n), c(n + 1, 0) {} - void update(int x, int delta) { - x += 1e5 + 1; - while (x <= n) { - c[x] += delta; - x += lowbit(x); + void update(int x, int v) { + for (; x <= n; x += x & -x) { + c[x] += v; } } int query(int x) { - x += 1e5 + 1; int s = 0; - while (x > 0) { + for (; x > 0; x -= x & -x) { s += c[x]; - x -= lowbit(x); } return s; } - - int lowbit(int x) { - x += 1e5 + 1; - return x & -x; - } }; class Solution { public: int subarraysWithMoreZerosThanOnes(vector& nums) { int n = nums.size(); - vector s(n + 1); - for (int i = 0; i < n; ++i) s[i + 1] = s[i] + (nums[i] == 1 ? 1 : -1); - BinaryIndexedTree* tree = new BinaryIndexedTree(n + 1); - int ans = 0; - const int MOD = 1e9 + 7; - for (int v : s) { - ans = (ans + tree->query(v - 1)) % MOD; - tree->update(v, 1); + 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; } @@ -232,51 +211,85 @@ type BinaryIndexedTree struct { } func newBinaryIndexedTree(n int) *BinaryIndexedTree { - n += 1e5 + 1 - c := make([]int, n+1) - return &BinaryIndexedTree{n, c} -} - -func (this *BinaryIndexedTree) lowbit(x int) int { - x += 1e5 + 1 - return x & -x + return &BinaryIndexedTree{n: n, c: make([]int, n+1)} } -func (this *BinaryIndexedTree) update(x, delta int) { - x += 1e5 + 1 - for x <= this.n { - this.c[x] += delta - x += this.lowbit(x) +func (bit *BinaryIndexedTree) update(x, v int) { + for ; x <= bit.n; x += x & -x { + bit.c[x] += v } } -func (this *BinaryIndexedTree) query(x int) int { - s := 0 - x += 1e5 + 1 - for x > 0 { - s += this.c[x] - x -= this.lowbit(x) +func (bit *BinaryIndexedTree) query(x int) (s int) { + for ; x > 0; x -= x & -x { + s += bit.c[x] } - return s + return } -func subarraysWithMoreZerosThanOnes(nums []int) int { +func subarraysWithMoreZerosThanOnes(nums []int) (ans int) { n := len(nums) - s := make([]int, n+1) - for i, v := range nums { - if v == 0 { - v = -1 + 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++ } - s[i+1] = s[i] + v + ans += tree.query(s - 1 + base) + ans %= mod + tree.update(s+base, 1) } - tree := newBinaryIndexedTree(n + 1) - ans := 0 - mod := int(1e9 + 7) - for _, v := range s { - ans = (ans + tree.query(v-1)) % mod - tree.update(v, 1) - } - return ans + return +} +``` + +### **TypeScript** + +```ts +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 === 0 ? -1 : 1; + ans += tree.query(s - 1 + base); + ans %= mod; + tree.update(s + base, 1); + } + return ans; } ``` diff --git a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README_EN.md b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README_EN.md index 20531468b5c48..9dbc065077bd9 100644 --- a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README_EN.md +++ b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README_EN.md @@ -50,7 +50,17 @@ The subarrays of size 1 that have more ones than zeros are: [1] ## Solutions -Binary Indexed Tree. +**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$. @@ -58,43 +68,38 @@ Binary Indexed Tree. ```python class BinaryIndexedTree: - def __init__(self, n): - n += int(1e5 + 1) + __slots__ = ["n", "c"] + + def __init__(self, n: int): self.n = n self.c = [0] * (n + 1) - @staticmethod - def lowbit(x): - x += int(1e5 + 1) - return x & -x - - def update(self, x, delta): - x += int(1e5 + 1) + def update(self, x: int, v: int): while x <= self.n: - self.c[x] += delta - x += BinaryIndexedTree.lowbit(x) + self.c[x] += v + x += x & -x - def query(self, x): - x += int(1e5 + 1) + def query(self, x: int) -> int: s = 0 - while x > 0: + while x: s += self.c[x] - x -= BinaryIndexedTree.lowbit(x) + x -= x & -x return s class Solution: def subarraysWithMoreZerosThanOnes(self, nums: List[int]) -> int: n = len(nums) - s = [0] - for v in nums: - s.append(s[-1] + (v or -1)) - tree = BinaryIndexedTree(n + 1) - MOD = int(1e9 + 7) - ans = 0 - for v in s: - ans = (ans + tree.query(v - 1)) % MOD - tree.update(v, 1) + base = n + 1 + tree = BinaryIndexedTree(n + base) + tree.update(base, 1) + mod = 10**9 + 7 + ans = s = 0 + for x in nums: + s += 1 if x else -1 + ans += tree.query(s - 1 + base) + ans %= mod + tree.update(s + base, 1) return ans ``` @@ -106,49 +111,38 @@ class BinaryIndexedTree { private int[] c; public BinaryIndexedTree(int n) { - n += (int) 1e5 + 1; this.n = n; c = new int[n + 1]; } - public void update(int x, int delta) { - x += (int) 1e5 + 1; - while (x <= n) { - c[x] += delta; - x += lowbit(x); + public void update(int x, int v) { + for (; x <= n; x += x & -x) { + c[x] += v; } } public int query(int x) { - x += (int) 1e5 + 1; int s = 0; - while (x > 0) { + for (; x > 0; x -= x & -x) { s += c[x]; - x -= lowbit(x); } return s; } - - public static int lowbit(int x) { - x += (int) 1e5 + 1; - return x & -x; - } } class Solution { - private static final int MOD = (int) 1e9 + 7; - public int subarraysWithMoreZerosThanOnes(int[] nums) { int n = nums.length; - int[] s = new int[n + 1]; - for (int i = 0; i < n; ++i) { - s[i + 1] = s[i] + (nums[i] == 1 ? 1 : -1); - } - BinaryIndexedTree tree = new BinaryIndexedTree(n + 1); - int ans = 0; - for (int v : s) { - ans = (ans + tree.query(v - 1)) % MOD; - tree.update(v, 1); + 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; } @@ -159,50 +153,42 @@ class Solution { ```cpp class BinaryIndexedTree { -public: +private: int n; vector c; - BinaryIndexedTree(int _n) - : n(_n + 1e5 + 1) - , c(_n + 1 + 1e5 + 1) {} +public: + BinaryIndexedTree(int n) : n(n), c(n + 1, 0) {} - void update(int x, int delta) { - x += 1e5 + 1; - while (x <= n) { - c[x] += delta; - x += lowbit(x); + void update(int x, int v) { + for (; x <= n; x += x & -x) { + c[x] += v; } } int query(int x) { - x += 1e5 + 1; int s = 0; - while (x > 0) { + for (; x > 0; x -= x & -x) { s += c[x]; - x -= lowbit(x); } return s; } - - int lowbit(int x) { - x += 1e5 + 1; - return x & -x; - } }; class Solution { public: int subarraysWithMoreZerosThanOnes(vector& nums) { int n = nums.size(); - vector s(n + 1); - for (int i = 0; i < n; ++i) s[i + 1] = s[i] + (nums[i] == 1 ? 1 : -1); - BinaryIndexedTree* tree = new BinaryIndexedTree(n + 1); - int ans = 0; - const int MOD = 1e9 + 7; - for (int v : s) { - ans = (ans + tree->query(v - 1)) % MOD; - tree->update(v, 1); + 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; } @@ -218,51 +204,85 @@ type BinaryIndexedTree struct { } func newBinaryIndexedTree(n int) *BinaryIndexedTree { - n += 1e5 + 1 - c := make([]int, n+1) - return &BinaryIndexedTree{n, c} + return &BinaryIndexedTree{n: n, c: make([]int, n+1)} } -func (this *BinaryIndexedTree) lowbit(x int) int { - x += 1e5 + 1 - return x & -x -} - -func (this *BinaryIndexedTree) update(x, delta int) { - x += 1e5 + 1 - for x <= this.n { - this.c[x] += delta - x += this.lowbit(x) +func (bit *BinaryIndexedTree) update(x, v int) { + for ; x <= bit.n; x += x & -x { + bit.c[x] += v } } -func (this *BinaryIndexedTree) query(x int) int { - s := 0 - x += 1e5 + 1 - for x > 0 { - s += this.c[x] - x -= this.lowbit(x) +func (bit *BinaryIndexedTree) query(x int) (s int) { + for ; x > 0; x -= x & -x { + s += bit.c[x] } - return s + return } -func subarraysWithMoreZerosThanOnes(nums []int) int { +func subarraysWithMoreZerosThanOnes(nums []int) (ans int) { n := len(nums) - s := make([]int, n+1) - for i, v := range nums { - if v == 0 { - v = -1 + 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++ } - s[i+1] = s[i] + v + ans += tree.query(s - 1 + base) + ans %= mod + tree.update(s+base, 1) } - tree := newBinaryIndexedTree(n + 1) - ans := 0 - mod := int(1e9 + 7) - for _, v := range s { - ans = (ans + tree.query(v-1)) % mod - tree.update(v, 1) - } - return ans + return +} +``` + +### **TypeScript** + +```ts +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 === 0 ? -1 : 1; + ans += tree.query(s - 1 + base); + ans %= mod; + tree.update(s + base, 1); + } + return ans; } ``` diff --git a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.cpp b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.cpp index 1abee0a5e40a3..5611786ace631 100644 --- a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.cpp +++ b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.cpp @@ -1,48 +1,40 @@ class BinaryIndexedTree { -public: +private: int n; vector c; - BinaryIndexedTree(int _n) - : n(_n + 1e5 + 1) - , c(_n + 1 + 1e5 + 1) {} +public: + BinaryIndexedTree(int n) : n(n), c(n + 1, 0) {} - void update(int x, int delta) { - x += 1e5 + 1; - while (x <= n) { - c[x] += delta; - x += lowbit(x); + void update(int x, int v) { + for (; x <= n; x += x & -x) { + c[x] += v; } } int query(int x) { - x += 1e5 + 1; int s = 0; - while (x > 0) { + for (; x > 0; x -= x & -x) { s += c[x]; - x -= lowbit(x); } return s; } - - int lowbit(int x) { - x += 1e5 + 1; - return x & -x; - } }; class Solution { public: int subarraysWithMoreZerosThanOnes(vector& nums) { int n = nums.size(); - vector s(n + 1); - for (int i = 0; i < n; ++i) s[i + 1] = s[i] + (nums[i] == 1 ? 1 : -1); - BinaryIndexedTree* tree = new BinaryIndexedTree(n + 1); - int ans = 0; - const int MOD = 1e9 + 7; - for (int v : s) { - ans = (ans + tree->query(v - 1)) % MOD; - tree->update(v, 1); + 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; } diff --git a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.go b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.go index 21a5e044f8fde..0e31abf996721 100644 --- a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.go +++ b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.go @@ -4,49 +4,38 @@ type BinaryIndexedTree struct { } func newBinaryIndexedTree(n int) *BinaryIndexedTree { - n += 1e5 + 1 - c := make([]int, n+1) - return &BinaryIndexedTree{n, c} + return &BinaryIndexedTree{n: n, c: make([]int, n+1)} } -func (this *BinaryIndexedTree) lowbit(x int) int { - x += 1e5 + 1 - return x & -x -} - -func (this *BinaryIndexedTree) update(x, delta int) { - x += 1e5 + 1 - for x <= this.n { - this.c[x] += delta - x += this.lowbit(x) +func (bit *BinaryIndexedTree) update(x, v int) { + for ; x <= bit.n; x += x & -x { + bit.c[x] += v } } -func (this *BinaryIndexedTree) query(x int) int { - s := 0 - x += 1e5 + 1 - for x > 0 { - s += this.c[x] - x -= this.lowbit(x) +func (bit *BinaryIndexedTree) query(x int) (s int) { + for ; x > 0; x -= x & -x { + s += bit.c[x] } - return s + return } -func subarraysWithMoreZerosThanOnes(nums []int) int { +func subarraysWithMoreZerosThanOnes(nums []int) (ans int) { n := len(nums) - s := make([]int, n+1) - for i, v := range nums { - if v == 0 { - v = -1 + 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++ } - s[i+1] = s[i] + v - } - tree := newBinaryIndexedTree(n + 1) - ans := 0 - mod := int(1e9 + 7) - for _, v := range s { - ans = (ans + tree.query(v-1)) % mod - tree.update(v, 1) + ans += tree.query(s - 1 + base) + ans %= mod + tree.update(s+base, 1) } - return ans + return } \ No newline at end of file diff --git a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.java b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.java index c895c6c9dac3e..bcd4b546e12c7 100644 --- a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.java +++ b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.java @@ -3,49 +3,38 @@ class BinaryIndexedTree { private int[] c; public BinaryIndexedTree(int n) { - n += (int) 1e5 + 1; this.n = n; c = new int[n + 1]; } - public void update(int x, int delta) { - x += (int) 1e5 + 1; - while (x <= n) { - c[x] += delta; - x += lowbit(x); + public void update(int x, int v) { + for (; x <= n; x += x & -x) { + c[x] += v; } } public int query(int x) { - x += (int) 1e5 + 1; int s = 0; - while (x > 0) { + for (; x > 0; x -= x & -x) { s += c[x]; - x -= lowbit(x); } return s; } - - public static int lowbit(int x) { - x += (int) 1e5 + 1; - return x & -x; - } } class Solution { - private static final int MOD = (int) 1e9 + 7; - public int subarraysWithMoreZerosThanOnes(int[] nums) { int n = nums.length; - int[] s = new int[n + 1]; - for (int i = 0; i < n; ++i) { - s[i + 1] = s[i] + (nums[i] == 1 ? 1 : -1); - } - BinaryIndexedTree tree = new BinaryIndexedTree(n + 1); - int ans = 0; - for (int v : s) { - ans = (ans + tree.query(v - 1)) % MOD; - tree.update(v, 1); + 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; } diff --git a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.py b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.py index 5a83fde8ff783..70ca9bf096056 100644 --- a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.py +++ b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.py @@ -1,39 +1,34 @@ class BinaryIndexedTree: - def __init__(self, n): - n += int(1e5 + 1) + __slots__ = ["n", "c"] + + def __init__(self, n: int): self.n = n self.c = [0] * (n + 1) - @staticmethod - def lowbit(x): - x += int(1e5 + 1) - return x & -x - - def update(self, x, delta): - x += int(1e5 + 1) + def update(self, x: int, v: int): while x <= self.n: - self.c[x] += delta - x += BinaryIndexedTree.lowbit(x) + self.c[x] += v + x += x & -x - def query(self, x): - x += int(1e5 + 1) + def query(self, x: int) -> int: s = 0 - while x > 0: + while x: s += self.c[x] - x -= BinaryIndexedTree.lowbit(x) + x -= x & -x return s class Solution: def subarraysWithMoreZerosThanOnes(self, nums: List[int]) -> int: n = len(nums) - s = [0] - for v in nums: - s.append(s[-1] + (v or -1)) - tree = BinaryIndexedTree(n + 1) - MOD = int(1e9 + 7) - ans = 0 - for v in s: - ans = (ans + tree.query(v - 1)) % MOD - tree.update(v, 1) + base = n + 1 + tree = BinaryIndexedTree(n + base) + tree.update(base, 1) + mod = 10**9 + 7 + ans = s = 0 + for x in nums: + s += 1 if x else -1 + ans += tree.query(s - 1 + base) + ans %= mod + tree.update(s + base, 1) return ans diff --git a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.ts b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.ts new file mode 100644 index 0000000000000..0ce68528e2c2b --- /dev/null +++ b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.ts @@ -0,0 +1,40 @@ +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 === 0 ? -1 : 1; + ans += tree.query(s - 1 + base); + ans %= mod; + tree.update(s + base, 1); + } + return ans; +} From 11861499fe0f2cb0bf593d42b9136d6d868569e5 Mon Sep 17 00:00:00 2001 From: yanglbme Date: Wed, 20 Dec 2023 09:53:06 +0800 Subject: [PATCH 2/2] fix: cpp code --- .../2031.Count Subarrays With More Ones Than Zeros/README.md | 4 +++- .../README_EN.md | 4 +++- .../Solution.cpp | 4 +++- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README.md b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README.md index 7fb63bf340ea0..736c0967fb3f6 100644 --- a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README.md +++ b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README.md @@ -165,7 +165,9 @@ private: vector c; public: - BinaryIndexedTree(int n) : n(n), c(n + 1, 0) {} + BinaryIndexedTree(int n) + : n(n) + , c(n + 1, 0) {} void update(int x, int v) { for (; x <= n; x += x & -x) { diff --git a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README_EN.md b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README_EN.md index 9dbc065077bd9..b93303396d396 100644 --- a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README_EN.md +++ b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/README_EN.md @@ -158,7 +158,9 @@ private: vector c; public: - BinaryIndexedTree(int n) : n(n), c(n + 1, 0) {} + BinaryIndexedTree(int n) + : n(n) + , c(n + 1, 0) {} void update(int x, int v) { for (; x <= n; x += x & -x) { diff --git a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.cpp b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.cpp index 5611786ace631..a9fe2dee03a78 100644 --- a/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.cpp +++ b/solution/2000-2099/2031.Count Subarrays With More Ones Than Zeros/Solution.cpp @@ -4,7 +4,9 @@ class BinaryIndexedTree { vector c; public: - BinaryIndexedTree(int n) : n(n), c(n + 1, 0) {} + BinaryIndexedTree(int n) + : n(n) + , c(n + 1, 0) {} void update(int x, int v) { for (; x <= n; x += x & -x) {