力扣刷题——3261. 统计满足 K 约束的子字符串数量 II

看了题目的两个初始用例，感觉能用前缀和和滑动窗口来解决，前缀和设定为从下标0到当前位置所有符合条件的答案数量，于是先写了一个：

vector<long long> countKConstraintSubstrings(string s, int k, vector<vector<int>> &queries)
{
    int n = s.size();
    vector<long long> prefix(n + 1, 0);
    vector<int> count(2, 0);
    int temLeft = 0;
    for (int i = 0; i < s.size(); i++)
    {
        count[s[i] - '0']++;
        while (count[0] > k && count[1] > k)
        {
            count[s[temLeft] - '0']--;
            temLeft++;
        }
        prefix[i + 1] = prefix[i] + i - temLeft + 1;
    }

    vector<long long> res;
    for (int i = 0; i < queries.size(); i++)
    {
        int left = queries[i][0];
        int right = queries[i][1];
        if (left)
            res.emplace_back(prefix[right + 1] - prefix[left + 1]);
        else
            res.emplace_back(prefix[right + 1] - prefix[left + 1] + 1);
    }
    return res;
}

但是这样存在一个问题，当用例为s = "010101", k = 2时就会开始报错。在query = [[0,4],[1,4]]时，能够得到对于[0,4]用例的正确答案15，而在计算[1,4]时使用前缀和相减的方式15 - 3，得到了错误的答案12，这是因为在这两个查询区域之中，任取用例都是合法的，在这种情况下不能够用前缀和的方式求解，而应该用(query[1] - query[0] + 1) * (query[1] - query[0]) / 2。此外也需要注意到，每个查询区域中的子数组可能不全是合法的，所以最终的答案应该是一个查询区域中，两种可能的答案形式之和，即全合法区域和非全合法区域，因此需要一个能够将两个区域分开的下标mid，这样每个查询的答案就为,非全合法部分prefix[query[1]] - prefix[mid]，与全合法部分(mid - query[0] + 1) * (mid - query[0]) / 2。
为了能够得到下标mid,在处理计算前缀和的过程中，需要保存对于每一次当前位置的一个合法区域，其中的左边界。于是有如下实现：

vector<long long> countKConstraintSubstrings(string s, int k, vector<vector<int>> &queries)
{
    int n = s.size();
    //这个实现跟上述文字描述有小出入，为了更优雅的写法，将前缀和的数组大小设置为n + 1
    vector<long long> prefix(n + 1, 0);
    vector<int> vecLeft(n, 0);
    vector<int> count(2, 0);
    int temLeft = 0;
    for (int i = 0; i < s.size(); i++)
    {
        count[s[i] - '0']++;
        while (count[0] > k && count[1] > k)
        {
            count[s[temLeft] - '0']--;
            temLeft++;
        }
        prefix[i + 1] = prefix[i] + i - temLeft + 1;
        vecLeft[i] = temLeft;
    }

    vector<long long> res;
    for (int i = 0; i < queries.size(); i++)
    {
        int left = queries[i][0];
        int right = queries[i][1];
        //因为在这个数组中，值是单调的，因此可以使用二分查找提速
        int mid = lower_bound(vecLeft.begin() + left, vecLeft.begin() + right + 1, left) - vecLeft.begin();
        res.emplace_back((long long)(mid - left) * (mid - left + 1) / 2 + prefix[right + 1] - prefix[mid]);
    }
    return res;
}