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

2024-11-13

题目：

给你一个 二进制 字符串 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
  • 所有查询互不相同

思路：

image-20241113092322599

代码：

class Solution {
public:
    vector<long long> countKConstraintSubstrings(string s, int k, vector<vector<int>>& queries) {
        int n = s.length();
        vector<int> left(n);
        vector<long long> sum(n + 1);
        int cnt[2]{}, l = 0;
        for (int i = 0; i < n; i++) {
            cnt[s[i] & 1]++;
            while (cnt[0] > k && cnt[1] > k) {
                cnt[s[l++] & 1]--;
            }
            left[i] = l; // 记录合法子串右端点 i 对应的最小左端点 l
            // 计算 i-left[i]+1 的前缀和
            sum[i + 1] = sum[i] + i - l + 1;
        }

        vector<long long> ans(queries.size());
        for (int i = 0; i < queries.size(); i++) {
            int l = queries[i][0], r = queries[i][1];
            // 如果区间内所有数都小于 l，结果是 j=r+1
            int j = lower_bound(left.begin() + l, left.begin() + r + 1, l) - left.begin();
            ans[i] = sum[r + 1] - sum[j] + (long long) (j - l + 1) * (j - l) / 2;
        }
        return ans;
    }
};