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

2024-11-12

题目:

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

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

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

返回一个整数,表示 s 的所有满足 k 约束 的子字符串的数量。

示例 1:

输入:s = “10101”, k = 1

输出:12

解释:

s 的所有子字符串中,除了 "1010""10101""0101" 外,其余子字符串都满足 k 约束。

示例 2:

输入:s = “1010101”, k = 2

输出:25

解释:

s 的所有子字符串中,除了长度大于 5 的子字符串外,其余子字符串都满足 k 约束。

示例 3:

输入:s = “11111”, k = 1

输出:15

解释:

s 的所有子字符串都满足 k 约束。

提示:

  • 1 <= s.length <= 50
  • 1 <= k <= s.length
  • s[i]'0''1'

思路:

计算以 0 为右端点的合法子串个数,以 1 为右端点的合法子串个数,……,以 n−1 为右端点的合法子串个数。

我们需要知道以 i 为右端点的合法子串,其左端点最小是多少。

由于随着 i 的变大,窗口内的字符数量变多,越不能满足题目要求,所以最小左端点会随着 i 的增大而增大,有单调性,因此可以用 滑动窗口 计算。

设以 i 为右端点的合法子串,其左端点最小是 left_i 。那么以 i 为右端点的合法子串,其左端点可以是 left_i , left_i + 1, …, i,一共 i − left_i + 1个,累加到答案中。

细节:字符 0 的 ASCII 值是偶数,字符 1 的 ASCII 值是奇数,所以可以用 ASCII 值 c mod 2 得到对应的数字。这也等价于和 1 计算 AND。

代码:

class Solution {
public:
    int countKConstraintSubstrings(string s, int k) {
        int ans = 0, left = 0;
        int cntZero = 0, cntOne = 0;
        for (int i = 0; i < s.size(); i++) {
            if (s[i] == '0') cntZero++;
            else cntOne++;
            while (cntOne > k && cntZero > k) {
                if (s[left] == '0') cntZero--;
                else cntOne--;
                left++;
            }
            ans += i - left + 1;
        }
        return ans;
    }
};