跳转至

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 <= 105
  • s[i] 是 '0' 或 '1'
  • 1 <= k <= s.length
  • 1 <= queries.length <= 105
  • queries[i] == [li, ri]
  • 0 <= li <= ri < s.length
  • 所有查询互不相同

解法

方法一：滑动窗口 + 前缀和

我们用两个变量 $\textit{cnt0}$ 和 $\textit{cnt1}$ 分别记录当前窗口内的 $0$ 和 $1$ 的个数，指针 $i$ 和 $j$ 分别标识窗口的左右边界。用一个数组 $d$ 记录每个位置 $i$ 右边第一个不满足 $k$ 约束的位置，初始时 $d[i] = n$。另外，用一个长度为 $n + 1$ 的前缀和数组 $\textit{pre}[i]$ 记录以前 $i$ 个位置作为右边界的满足 $k$ 约束的子字符串的个数。

当我们右移窗口时，如果窗口内的 $0$ 和 $1$ 的个数都大于 $k$，我们将 $d[i]$ 更新为 $j$，表示位置 $i$ 右边第一个不满足 $k$ 约束的位置。然后我们将 $i$ 右移一位，直到窗口内的 $0$ 和 $1$ 的个数都不大于 $k$。此时，我们可以计算出以 $j$ 为右边界的满足 $k$ 约束的子字符串的个数，即 $j - i + 1$，我们更新到前缀和数组中。

最后，对于每个查询 $[l, r]$，我们首先找出 $l$ 右边第一个不满足 $k$ 约束的位置 $p$，那么 $p = \min(r + 1, d[l])$，那么 $[l, p - 1]$ 的所有子字符串都满足 $k$ 约束，个数为 $(1 + p - l) \times (p - l) / 2$，然后，我们计算以 $[p, r]$ 为右边界的满足 $k$ 约束的子字符串的个数，即 $\textit{pre}[r + 1] - \textit{pre}[p]$，最后将两者相加即可。

时间复杂度 $O(n + m)$，空间复杂度 $O(n)$。其中 $n$ 和 $m$ 分别为字符串 $s$ 的长度和查询数组 $\textit{queries}$ 的长度。

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution:
    def countKConstraintSubstrings(
        self, s: str, k: int, queries: List[List[int]]
    ) -> List[int]:
        cnt = [0, 0]
        i, n = 0, len(s)
        d = [n] * n
        pre = [0] * (n + 1)
        for j, x in enumerate(map(int, s)):
            cnt[x] += 1
            while cnt[0] > k and cnt[1] > k:
                d[i] = j
                cnt[int(s[i])] -= 1
                i += 1
            pre[j + 1] = pre[j] + j - i + 1
        ans = []
        for l, r in queries:
            p = min(r + 1, d[l])
            a = (1 + p - l) * (p - l) // 2
            b = pre[r + 1] - pre[p]
            ans.append(a + b)
        return ans

 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
class Solution {
    public long[] countKConstraintSubstrings(String s, int k, int[][] queries) {
        int[] cnt = new int[2];
        int n = s.length();
        int[] d = new int[n];
        Arrays.fill(d, n);
        long[] pre = new long[n + 1];
        for (int i = 0, j = 0; j < n; ++j) {
            cnt[s.charAt(j) - '0']++;
            while (cnt[0] > k && cnt[1] > k) {
                d[i] = j;
                cnt[s.charAt(i++) - '0']--;
            }
            pre[j + 1] = pre[j] + j - i + 1;
        }
        int m = queries.length;
        long[] ans = new long[m];
        for (int i = 0; i < m; ++i) {
            int l = queries[i][0], r = queries[i][1];
            int p = Math.min(r + 1, d[l]);
            long a = (1L + p - l) * (p - l) / 2;
            long b = pre[r + 1] - pre[p];
            ans[i] = a + b;
        }
        return ans;
    }
}

 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
class Solution {
public:
    vector<long long> countKConstraintSubstrings(string s, int k, vector<vector<int>>& queries) {
        int cnt[2]{};
        int n = s.size();
        vector<int> d(n, n);
        long long pre[n + 1];
        pre[0] = 0;
        for (int i = 0, j = 0; j < n; ++j) {
            cnt[s[j] - '0']++;
            while (cnt[0] > k && cnt[1] > k) {
                d[i] = j;
                cnt[s[i++] - '0']--;
            }
            pre[j + 1] = pre[j] + j - i + 1;
        }
        vector<long long> ans;
        for (const auto& q : queries) {
            int l = q[0], r = q[1];
            int p = min(r + 1, d[l]);
            long long a = (1LL + p - l) * (p - l) / 2;
            long long b = pre[r + 1] - pre[p];
            ans.push_back(a + b);
        }
        return ans;
    }
};

 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
func countKConstraintSubstrings(s string, k int, queries [][]int) (ans []int64) {
    cnt := [2]int{}
    n := len(s)
    d := make([]int, n)
    for i := range d {
        d[i] = n
    }
    pre := make([]int, n+1)
    for i, j := 0, 0; j < n; j++ {
        cnt[s[j]-'0']++
        for cnt[0] > k && cnt[1] > k {
            d[i] = j
            cnt[s[i]-'0']--
            i++
        }
        pre[j+1] = pre[j] + j - i + 1
    }
    for _, q := range queries {
        l, r := q[0], q[1]
        p := min(r+1, d[l])
        a := (1 + p - l) * (p - l) / 2
        b := pre[r+1] - pre[p]
        ans = append(ans, int64(a+b))
    }
    return
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
function countKConstraintSubstrings(s: string, k: number, queries: number[][]): number[] {
    const cnt: [number, number] = [0, 0];
    const n = s.length;
    const d: number[] = Array(n).fill(n);
    const pre: number[] = Array(n + 1).fill(0);
    for (let i = 0, j = 0; j < n; ++j) {
        cnt[+s[j]]++;
        while (Math.min(cnt[0], cnt[1]) > k) {
            d[i] = j;
            cnt[+s[i++]]--;
        }
        pre[j + 1] = pre[j] + j - i + 1;
    }
    const ans: number[] = [];
    for (const [l, r] of queries) {
        const p = Math.min(r + 1, d[l]);
        const a = ((1 + p - l) * (p - l)) / 2;
        const b = pre[r + 1] - pre[p];
        ans.push(a + b);
    }
    return ans;
}

评论