3261.统计满足K约束的子字符串数量II

目标

给你一个 二进制 字符串 s 和一个整数 k。

另给你一个二维整数数组 queries ,其中 queries[i] = [li, ri] 。

如果一个 二进制字符串 满足以下任一条件,则认为该字符串满足 k 约束:

  • 字符串中 0 的数量最多为 k。
  • 字符串中 1 的数量最多为 k。

返回一个整数数组 answer ,其中 answer[i] 表示 s[li..ri] 中满足 k 约束 的 子字符串 的数量。

示例 1:

输入:s = "0001111", k = 2, queries = [[0,6]]
输出:[26]
解释:
对于查询 [0, 6], s[0..6] = "0001111" 的所有子字符串中,除 s[0..5] = "000111" 和 s[0..6] = "0001111" 外,其余子字符串都满足 k 约束。

示例 2:

输入:s = "010101", k = 1, queries = [[0,5],[1,4],[2,3]]
输出:[15,9,3]
解释:
s 的所有子字符串中,长度大于 3 的子字符串都不满足 k 约束。

说明:

  • 1 <= s.length <= 10^5
  • s[i] 是 '0' 或 '1'
  • 1 <= k <= s.length
  • 1 <= queries.length <= 10^5
  • queries[i] == [li, ri]
  • 0 <= li <= ri < s.length
  • 所有查询互不相同

思路

定义二进制字符串满足 k 约束的条件是字符串中 0 的个数 或者 1 的个数都不超过 k。求给定字符串满足 k 约束的子字符串数量。子字符串 是字符串中 连续非空 字符序列。

这道题与昨天的 3258.统计满足K约束的子字符串数量I 相比多了一个查询数组,并且字符串的长度也来到了 10^5,返回结果是 long[],暴力枚举会超时。

滑动窗口的时间复杂度为 O(n),不可能对每一次查询都重新滑动一遍。显然需要在滑动的过程中记录下查询的结果。每次滑动我们可以得到一个区间 [l, r],这个区间的所有子数组是合法的。

使用一维数组记录这个区间,使用下标与值的映射,我们有两种选择:

  • 记录的左端点的最大右端点,即 left[l] = r
  • 记录的右端点的最小左端点,即 right[r] = l

对于查询区间 [ql, qr],它与我们已知的合法区间存在两种关系,被包含或者部分相交。

  • 如果是被包含的关系,那么查询区间的所有子数组均合法,子数组个数为 (qr - ql + 1) * (qr - ql + 2) / 2
  • 如果是相交的关系,说明 ql < r[ql, r] 的所有子数组是合法的,剩下的 [r + 1, qr] 的合法子数组如何求?可以在滑动过程中记录以每个元素为终点的合法子数组个数,以前缀和的形式保存。

这里区间的划分与结果集的构成是非常有讲究的。前缀和记录的是以元素为 终点 的合法子数组,如果我们在滑动的过程中根据查询区间终点匹配当前元素,即 qr == r,那么可能的情况为 ql >= l 查询区间被完全包含。如果 ql < l 则查询区间与合法区间相交。如果这时使用前缀和计算 [ql, l],使用公式计算 [l, r] 就错了。因为后面区间使用公式计算的子数组并不包括以前面区间内的元素为起点的子数组,并且前缀和记录的子数组的起点可能在查询的左边界之外。而反过来前面使用公式计算,后面使用前缀和计算,被减去的那部分子数组个数中就包含了以 前面区间所有元素 为终点的子数组,也就是前面使用公式计算的子数组个数。我们不用担心后面通过前缀和计算的子数组的起点超出左边界,如果超出的话,一定是被包含的关系。

核心点在于想清楚这两部分集合的关系, [i, j] 的所有子数组包括了以 b ∈ [i, j] 为终点,a ∈ [i, b] 为起点的子数组。而使用前缀和相减计算的是所有以 c ∈ [m, n] 为终点的合法子数组,起点可以不在该区间,但是不会超出 ql

代码


/**
 * @date 2024-11-13 0:25
 */
public class CountKConstraintSubstrings3261 {

    public long[] countKConstraintSubstrings_v1(String s, int k, int[][] queries) {
        int n = s.length();
        char[] chars = s.toCharArray();
        int[] cnt = new int[2];
        long[] res = new long[queries.length];
        long[] prefix = new long[n + 1];
        int[] right = new int[n];
        Arrays.fill(right, n);
        int l = 0;
        for (int i = 0; i < n; i++) {
            cnt[chars[i] - '0']++;
            while (l <= i && cnt[0] > k && cnt[1] > k) {
                right[l] = i;
                cnt[chars[l++] - '0']--;
            }
            prefix[i + 1] += prefix[i] + i - l + 1;
        }
        for (int i = 0; i < queries.length; i++) {
            int ql = queries[i][0];
            int qr = queries[i][1];
            int r = Math.min(right[ql], qr + 1);
            res[i] = (r - ql + 1L) * (r - ql) / 2 + prefix[qr + 1] - prefix[r];

        }
        return res;
    }

}

性能