3261. Count Substrings That Satisfy K-Constraint II
Description
You are given a binary string s and an integer k.
You are also given a 2D integer array queries, where queries[i] = [li, ri].
A binary string satisfies the k-constraint if either of the following conditions holds:
- The number of
0's in the string is at mostk. - The number of
1's in the string is at mostk.
Return an integer array answer, where answer[i] is the number of substrings of s[li..ri] that satisfy the k-constraint.
Example 1:
Input: s = "0001111", k = 2, queries = [[0,6]]
Output: [26]
Explanation:
For the query [0, 6], all substrings of s[0..6] = "0001111" satisfy the k-constraint except for the substrings s[0..5] = "000111" and s[0..6] = "0001111".
Example 2:
Input: s = "010101", k = 1, queries = [[0,5],[1,4],[2,3]]
Output: [15,9,3]
Explanation:
The substrings of s with a length greater than 3 do not satisfy the k-constraint.
Constraints:
1 <= s.length <= 105s[i]is either'0'or'1'.1 <= k <= s.length1 <= queries.length <= 105queries[i] == [li, ri]0 <= li <= ri < s.length- All queries are distinct.
Solutions
Solution 1: Sliding Window + Prefix Sum
We use two variables $\textit{cnt0}$ and $\textit{cnt1}$ to record the number of $0$s and $1$s in the current window, respectively. Pointers $i$ and $j$ mark the left and right boundaries of the window. We use an array $d$ to record the first position to the right of each position $i$ that does not satisfy the $k$ constraint, initially setting $d[i] = n$. Additionally, we use a prefix sum array $\textit{pre}[i]$ of length $n + 1$ to record the number of substrings that satisfy the $k$ constraint with the right boundary at position $i$.
When we move the window to the right, if the number of $0$s and $1$s in the window both exceed $k$, we update $d[i]$ to $j$, indicating that the first position to the right of $i$ that does not satisfy the $k$ constraint is $j$. Then we move $i$ one position to the right until the number of $0$s and $1$s in the window are both less than or equal to $k$. At this point, we can calculate the number of substrings that satisfy the $k$ constraint with the right boundary at $j$, which is $j - i + 1$, and update this in the prefix sum array.
Finally, for each query $[l, r]$, we first find the first position $p$ to the right of $l$ that does not satisfy the $k$ constraint, which is $p = \min(r + 1, d[l])$. All substrings in the range $[l, p - 1]$ satisfy the $k$ constraint, and the number of such substrings is $(1 + p - l) \times (p - l) / 2$. Then, we calculate the number of substrings that satisfy the $k$ constraint with the right boundary in the range $[p, r]$, which is $\textit{pre}[r + 1] - \textit{pre}[p]$. Finally, we add the two results together.
The time complexity is $O(n + m)$, and the space complexity is $O(n)$. Here, $n$ and $m$ are the lengths of the string $s$ and the query array $\textit{queries}$, respectively.
Python3
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
Java
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;
}
}
C++
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;
}
};
Go
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
}
TypeScript
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;
}