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 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.

 

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 <= 105
• s[i] is either '0' or '1'.
• 1 <= k <= s.length
• 1 <= queries.length <= 105
• queries[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.

Python3JavaC++GoTypeScript

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

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;
    }
}

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;
    }
};

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
}

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;
}

Comments