LC 2031. 1 比 0 多的子数组个数

题目描述

这是 LeetCode 上的 2031. 1 比 0 多的子数组个数 ，难度为 中等。

给你一个只包含 $0$ 和 $1$ 的数组 $nums$，请返回 $1$ 的数量 大于 $04 的数量的子数组的个数。

由于答案可能很大，请返回答案对 $10^9 + 7$ 取余 的结果。

一个 子数组 指的是原数组中连续的一个子序列。

示例 1:

1
2
3
4
5
6
7
8
9
10
输入: nums = [0,1,1,0,1]

输出: 9

解释:
长度为 1 的、1 的数量大于 0 的数量的子数组有: [1], [1], [1]
长度为 2 的、1 的数量大于 0 的数量的子数组有: [1,1]
长度为 3 的、1 的数量大于 0 的数量的子数组有: [0,1,1], [1,1,0], [1,0,1]
长度为 4 的、1 的数量大于 0 的数量的子数组有: [1,1,0,1]
长度为 5 的、1 的数量大于 0 的数量的子数组有: [0,1,1,0,1]

示例 2:
1
2
3
4
5
6
输入: nums = [0]

输出: 0

解释:
没有子数组的 1 的数量大于 0 的数量。

示例 3:
1
2
3
4
5
6
输入: nums = [1]

输出: 1

解释:
长度为 1 的、1 的数量大于 0 的数量的子数组有: [1]

提示:

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

树状数组

为了方便，我们调整数组 $nums$ 下标从 $1$ 开始。同时将 $nums[i] = 0$ 的值看作 $nums[i] = -1$，并预处理出前缀和数组 $sum$。对于任意 $sum[i]$ 而言，其值域范围为 $[-n, n]$，我们可以通过对 $sum[i]$ 做整体 $n + 1$ 的偏移，将值域映射到 $[1, 2 * n + 1]$。

对于任意一个子数组 $i…j$ 而言，如果满足 $1$ 的数量大于 $0$，则必然有 $sum[j] - sum[i - 1] > 0$。

因此在求解以 $nums[j]$ 为右端点的「满足 $1$ 数量大于 $0$ 数量」的子数组个数时，等价于在问 $[0, j - 1]$ 范围内有多少个下标 $i$ 满足 $sum[i] < sum[j]$（即有多少个下标可以作为左端点）。

求解比 $x$ 小的数有多少个，可以使用「树状数组」来做。

具体的，我们可以遍历每个位置，假设当前处理到的位置是 $i$，其前缀和值 $t = sum[i]$，我们可以通过查询小于等于 $t - 1$ 的值数量来得知以 $i$ 为右端点的合法子数组数量，统计后，我们需要对数值为 $t$ 出现次数进行加一。

代码：

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
class Solution {
    int N = 200010, MOD = (int)1e9+7;
    int[] tr = new int[N];
    int n;
    int lowbit(int x) {
        return x & -x;
    }
    void add(int x, int v) {
        for (int i = x; i <= 2 * n + 1; i += lowbit(i)) tr[i] += v;
    }
    int query(int x) {
        int ans = 0;
        for (int i = x; i > 0; i -= lowbit(i)) ans += tr[i];
        return ans;
    }
    public int subarraysWithMoreZerosThanOnes(int[] nums) {
        n = nums.length;
        int ans = 0;
        int[] sum = new int[n + 10];
        for (int i = 1; i <= n; i++) sum[i] = sum[i - 1] + (nums[i - 1] == 0 ? -1 : 1);
        for (int i = 0; i <= n; i++) {
            int t = sum[i] + n + 1;
            ans = (ans + query(t - 1)) % MOD;
            add(t, 1);
        }
        return ans;
    }
}

  • 时间复杂度：预处理前缀和数组的复杂度为 $O(n)$；查询和插入的复杂度均为 $O(\log{n})$，整体复杂度为 $O(n\log{n})$
  • 空间复杂度：$O(n)$

最后

这是我们「刷穿 LeetCode」系列文章的第 No.2031 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：https://github.com/SharingSource/LogicStack-LeetCode 。

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。