Problem

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 most k.
  • The number of 1’s in the string is at most k.

Return an integer array answer, where answer[i] is the number of substrings of s[li..ri] that satisfy the k-constraint.

Examples

Example 1

 1
 2
 3
 4
 5
 6
 7
 8
 9
10


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

1
2
3
4
5
6
7
8
9


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 <= 10^5
  • s[i] is either '0' or '1'.
  • 1 <= k <= s.length
  • 1 <= queries.length <= 10^5
  • queries[i] == [li, ri]
  • 0 <= li <= ri < s.length
  • All queries are distinct.

Solution

Method 1 – Prefix Sums and Binary Search

Intuition

For each query, we want to count the number of substrings in s[li..ri] that satisfy the k-constraint (at most k zeros or at most k ones). Since s can be large, we use prefix sums to quickly count zeros and ones in any substring, and for each starting index, use binary search to find the farthest valid end index.

Approach

  1. Precompute prefix sums for zeros and ones in s.
  2. For each query [l, r]:
    • For each start index i in [l, r]:
      • Use binary search to find the maximum j in [i, r] such that the substring s[i..j] has at most k zeros or at most k ones.
      • The number of valid substrings starting at i is (max_j - i + 1).
    • Sum over all i to get the answer for the query.
  3. Return the answer array.

Code

 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
28
29
30
31
32

class Solution {
public:
    vector<int> countSubstrings(string s, int k, vector<vector<int>>& queries) {
        int n = s.size();
        vector<int> p0(n+1), p1(n+1);
        for (int i = 0; i < n; ++i) {
            p0[i+1] = p0[i] + (s[i] == '0');
            p1[i+1] = p1[i] + (s[i] == '1');
        }
        vector<int> ans;
        for (auto& q : queries) {
            int l = q[0], r = q[1], res = 0;
            for (int i = l; i <= r; ++i) {
                int lo = i, hi = r, best = i-1;
                while (lo <= hi) {
                    int mid = (lo + hi) / 2;
                    int c0 = p0[mid+1] - p0[i];
                    int c1 = p1[mid+1] - p1[i];
                    if (c0 <= k || c1 <= k) {
                        best = mid;
                        lo = mid + 1;
                    } else {
                        hi = mid - 1;
                    }
                }
                res += best - i + 1;
            }
            ans.push_back(res);
        }
        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
28
29
30
31
32
33
34
35
36
37

type Solution struct{}
func (Solution) CountSubstrings(s string, k int, queries [][]int) []int {
    n := len(s)
    p0 := make([]int, n+1)
    p1 := make([]int, n+1)
    for i := 0; i < n; i++ {
        p0[i+1] = p0[i]
        p1[i+1] = p1[i]
        if s[i] == '0' {
            p0[i+1]++
        } else {
            p1[i+1]++
        }
    }
    ans := make([]int, len(queries))
    for qi, q := range queries {
        l, r := q[0], q[1]
        res := 0
        for i := l; i <= r; i++ {
            lo, hi, best := i, r, i-1
            for lo <= hi {
                mid := (lo + hi) / 2
                c0 := p0[mid+1] - p0[i]
                c1 := p1[mid+1] - p1[i]
                if c0 <= k || c1 <= k {
                    best = mid
                    lo = mid + 1
                } else {
                    hi = mid - 1
                }
            }
            res += best - i + 1
        }
        ans[qi] = res
    }
    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
28
29
30
31

class Solution {
    public int[] countSubstrings(String s, int k, int[][] queries) {
        int n = s.length();
        int[] p0 = new int[n+1], p1 = new int[n+1];
        for (int i = 0; i < n; ++i) {
            p0[i+1] = p0[i] + (s.charAt(i) == '0' ? 1 : 0);
            p1[i+1] = p1[i] + (s.charAt(i) == '1' ? 1 : 0);
        }
        int[] ans = new int[queries.length];
        for (int qi = 0; qi < queries.length; ++qi) {
            int l = queries[qi][0], r = queries[qi][1], res = 0;
            for (int i = l; i <= r; ++i) {
                int lo = i, hi = r, best = i-1;
                while (lo <= hi) {
                    int mid = (lo + hi) / 2;
                    int c0 = p0[mid+1] - p0[i];
                    int c1 = p1[mid+1] - p1[i];
                    if (c0 <= k || c1 <= k) {
                        best = mid;
                        lo = mid + 1;
                    } else {
                        hi = mid - 1;
                    }
                }
                res += best - i + 1;
            }
            ans[qi] = res;
        }
        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
28
29
30
31
32
33

class Solution {
    fun countSubstrings(s: String, k: Int, queries: Array<IntArray>): IntArray {
        val n = s.length
        val p0 = IntArray(n+1)
        val p1 = IntArray(n+1)
        for (i in 0 until n) {
            p0[i+1] = p0[i] + if (s[i] == '0') 1 else 0
            p1[i+1] = p1[i] + if (s[i] == '1') 1 else 0
        }
        val ans = IntArray(queries.size)
        for ((qi, q) in queries.withIndex()) {
            val l = q[0]; val r = q[1]
            var res = 0
            for (i in l..r) {
                var lo = i; var hi = r; var best = i-1
                while (lo <= hi) {
                    val mid = (lo + hi) / 2
                    val c0 = p0[mid+1] - p0[i]
                    val c1 = p1[mid+1] - p1[i]
                    if (c0 <= k || c1 <= k) {
                        best = mid
                        lo = mid + 1
                    } else {
                        hi = mid - 1
                    }
                }
                res += best - i + 1
            }
            ans[qi] = res
        }
        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

class Solution:
    def countSubstrings(self, s: str, k: int, queries: list[list[int]]) -> list[int]:
        n = len(s)
        p0 = [0] * (n+1)
        p1 = [0] * (n+1)
        for i in range(n):
            p0[i+1] = p0[i] + (s[i] == '0')
            p1[i+1] = p1[i] + (s[i] == '1')
        ans = []
        for l, r in queries:
            res = 0
            for i in range(l, r+1):
                lo, hi, best = i, r, i-1
                while lo <= hi:
                    mid = (lo + hi) // 2
                    c0 = p0[mid+1] - p0[i]
                    c1 = p1[mid+1] - p1[i]
                    if c0 <= k or c1 <= k:
                        best = mid
                        lo = mid + 1
                    else:
                        hi = mid - 1
                res += best - i + 1
            ans.append(res)
        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
28
29
30
31
32
33
34
35

impl Solution {
    pub fn count_substrings(s: String, k: i32, queries: Vec<Vec<i32>>) -> Vec<i32> {
        let n = s.len();
        let s = s.as_bytes();
        let mut p0 = vec![0; n+1];
        let mut p1 = vec![0; n+1];
        for i in 0..n {
            p0[i+1] = p0[i] + if s[i] == b'0' { 1 } else { 0 };
            p1[i+1] = p1[i] + if s[i] == b'1' { 1 } else { 0 };
        }
        let mut ans = Vec::new();
        for q in queries.iter() {
            let l = q[0] as usize;
            let r = q[1] as usize;
            let mut res = 0;
            for i in l..=r {
                let (mut lo, mut hi, mut best) = (i, r, i-1);
                while lo <= hi {
                    let mid = (lo + hi) / 2;
                    let c0 = p0[mid+1] - p0[i];
                    let c1 = p1[mid+1] - p1[i];
                    if c0 <= k || c1 <= k {
                        best = mid;
                        lo = mid + 1;
                    } else {
                        hi = mid - 1;
                    }
                }
                res += best as i32 - i as i32 + 1;
            }
            ans.push(res);
        }
        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
28
29
30
31
32

class Solution {
    countSubstrings(s: string, k: number, queries: number[][]): number[] {
        const n = s.length;
        const p0 = Array(n+1).fill(0);
        const p1 = Array(n+1).fill(0);
        for (let i = 0; i < n; ++i) {
            p0[i+1] = p0[i] + (s[i] === '0' ? 1 : 0);
            p1[i+1] = p1[i] + (s[i] === '1' ? 1 : 0);
        }
        const ans = [];
        for (const [l, r] of queries) {
            let res = 0;
            for (let i = l; i <= r; ++i) {
                let lo = i, hi = r, best = i-1;
                while (lo <= hi) {
                    const mid = (lo + hi) >> 1;
                    const c0 = p0[mid+1] - p0[i];
                    const c1 = p1[mid+1] - p1[i];
                    if (c0 <= k || c1 <= k) {
                        best = mid;
                        lo = mid + 1;
                    } else {
                        hi = mid - 1;
                    }
                }
                res += best - i + 1;
            }
            ans.push(res);
        }
        return ans;
    }
}

Complexity

  • ⏰ Time complexity: O(Q * N * log N), where Q is the number of queries and N is the length of s. For each query, for each start index, we binary search the end index.
  • 🧺 Space complexity: O(N) for the prefix sum arrays.