LeetCode Solutions Blog

Count Substrings That Satisfy K-Constraint II

Number: 3546

Difficulty: Hard

Paid? No

Companies: N/A

Problem Description

Given a binary string s and an integer k, a substring of s is said to satisfy the k-constraint if either the number of 0’s in the substring is at most k or the number of 1’s is at most k. You are also given an array of queries, where each query is a pair [l, r] indicating a segment of s. For each query, return the number of substrings (nonempty contiguous segments) within s[l..r] that satisfy the k-constraint.

Key Insights

Instead of iterating over all O(n²) substrings to check the constraint, notice that the property is monotonic: if a substring s[i..j] is valid then every shorter substring starting at i is also valid.
A substring is invalid only if both counts – the number of 0’s and 1’s – are greater than k. Equivalently, it is valid if at least one of these counts is at most k.
For each starting index i in s, we can compute the farthest ending index valid_end[i] for which every substring s[i..j] (for j from i up to valid_end[i]) satisfies the constraint. Once an index j causes both counts to exceed k, all longer substrings starting at i become invalid.
For any query [l, r] the count of valid substrings can be computed by summing for each starting position i from l to r the number: min(valid_end[i], r) - i + 1.
To answer many queries efficiently, we can precompute valid_end for all indices using a sliding window. Then, because each query [l, r] involves summing a function f(i, r) defined by a “min” over precomputed values, we can answer the queries offline using a Binary Indexed Tree (Fenwick tree). In particular, we split indices in [l, r] into two groups:
- Group 1: indices where valid_end[i] ≤ r. Here the contribution is F[i] = valid_end[i] − i + 1.
- Group 2: indices where valid_end[i] > r. For these indices the maximum substring ending is r, so contribution is r − i + 1.
Finally, by precomputing a prefix-sum array over the indices (to quickly get the sum of i’s over a given range) and by processing queries in order of increasing r, we can answer each query in O(log n) time.

Space and Time Complexity

Time Complexity: O((n + Q) log n), where n is the length of the string and Q is the number of queries. Space Complexity: O(n + Q), to store the valid_end array, prefix sums, and Fenwick trees.

Solution

We first precompute valid_end for every starting index i using a sliding-window approach. For each i, we move the right pointer until adding the next character would result in both the count of 0’s and 1’s exceeding k. This gives us valid_end[i] as the largest j such that every substring s[i..j] is valid.

Now, for a query [l, r], the contribution from each starting index i (l ≤ i ≤ r) is: if valid_end[i] ≤ r, then contribution = valid_end[i] − i + 1; if valid_end[i] > r, then contribution = r − i + 1. Thus the answer for the query is sum from i = l to r of min(valid_end[i], r) − i + 1.

Since we have up to 10⁵ queries, we process the queries offline (sorted by r) and use Fenwick trees to quickly sum contributions for indices satisfying valid_end[i] ≤ current r. In parallel, we use a precomputed prefix sum over indices to calculate the sum of (r − i + 1) over the remaining indices. Fenwick trees are updated in the order of increasing valid_end[i].

Below are code solutions in Python, JavaScript, C++, and Java with detailed inline comments.

Code Solutions

# Python solution
# We define a Fenwick Tree class with update and query support.
class FenwickTree:
    def __init__(self, size):
        self.n = size
        self.bit = [0] * (size + 1)
    
    # update index pos (0-indexed) with value delta
    def update(self, pos, delta):
        i = pos + 1
        while i <= self.n:
            self.bit[i] += delta
            i += i & -i
            
    # query sum from index 0 to pos (0-indexed)
    def query(self, pos):
        s = 0
        i = pos + 1
        while i:
            s += self.bit[i]
            i -= i & -i
        return s
    
    # query range [l, r] (0-indexed)
    def range_query(self, l, r):
        return self.query(r) - (self.query(l - 1) if l > 0 else 0)

def countValidSubstrings(s, k, queries):
    n = len(s)
    # Precompute valid_end for each index i using sliding window
    valid_end = [0] * n
    count0 = 0
    count1 = 0
    j = 0
    for i in range(n):
        # move j as far as possible while the window [i..j] is valid.
        while j < n:
            if s[j] == '0':
                temp0 = count0 + 1
                temp1 = count1
            else:
                temp0 = count0
                temp1 = count1 + 1
            # Check if adding s[j] would make both counts exceed k
            if temp0 > k and temp1 > k:
                break
            count0, count1 = temp0, temp1
            j += 1
        valid_end[i] = j - 1  # the last valid index for starting at i
        # Move the left pointer: remove s[i] from counts
        if i < j:
            if s[i] == '0':
                count0 -= 1
            else:
                count1 -= 1
        else:
            # If i caught up with j, then move j one step further.
            j = i + 1

    # Precompute an array F[i] = valid_end[i] - i + 1
    F = [valid_end[i] - i + 1 for i in range(n)]
    
    # Build prefix sum for indices for fast range sum queries.
    prefix = [0] * n
    prefix[0] = 0  # prefix sum of indices; note: value stored = index itself
    for i in range(1, n):
        prefix[i] = prefix[i - 1] + i

    # Process queries offline; each query: (r, l, original_index).
    offline_queries = []
    for idx, (l, r) in enumerate(queries):
        offline_queries.append((r, l, idx))
    offline_queries.sort()  # sort by r

    # Build Fenwick trees to support queries on indices.
    # ftF: sum of F[i] for indices added (those with valid_end[i] <= current r)
    ftF = FenwickTree(n)
    # ftCount: count of indices added.
    ftCount = FenwickTree(n)
    # ftIdx: sum of indices i for added indices.
    ftIdx = FenwickTree(n)
    
    # Sort indices by valid_end values.
    idx_by_valid = sorted(range(n), key=lambda i: valid_end[i])
    pos = 0   # pointer for idx_by_valid
    
    answers = [0] * len(queries)
    # Process each query in order of increasing r.
    for r, l, qidx in offline_queries:
        # Add all indices i which have valid_end[i] <= r.
        while pos < n and valid_end[idx_by_valid[pos]] <= r:
            i = idx_by_valid[pos]
            ftF.update(i, F[i])
            ftCount.update(i, 1)
            ftIdx.update(i, i)
            pos += 1
        # Count S1: For indices in [l, r] that satisfy valid_end[i] <= r.
        S1 = ftF.range_query(l, r)
        count1 = ftCount.range_query(l, r)
        sumIdx = ftIdx.range_query(l, r)
        # For indices where valid_end[i] > r, contribution = r - i + 1.
        totalIndices = (r - l + 1)
        count2 = totalIndices - count1
        # Sum of contributions for these indices.
        # Sum_{i in group2}(r - i + 1) = count2*(r+1) - (sum over i for group2)
        # To get sum over i for group2, we use precomputed prefix sums.
        totalSumIndices = prefix[r] - (prefix[l - 1] if l > 0 else 0)
        sumGroup2 = totalSumIndices - sumIdx
        contrib2 = count2 * (r + 1) - sumGroup2
        answers[qidx] = S1 + contrib2

    return answers

# Example usage:
s = "0001111"
k = 2
queries = [[0, 6]]
print(countValidSubstrings(s, k, queries))  # Expected output: [26]