LeetCode Problem 3261

Difficulty: 🔴 Hard
Topics: Array, String, Binary Search, Sliding Window, Prefix Sum

Solution

LeetCode 3261 - Count Substrings That Satisfy K-Constraint II

Problem Understanding

We are given a binary string s, an integer k, and many range queries.

For each query [l, r], we only consider the substring s[l..r]. Among all substrings completely contained inside this range, we must count how many satisfy the k-constraint.

A substring satisfies the constraint if at least one of the following is true:

• The number of 0s is at most k.
• The number of 1s is at most k.

Notice that a substring becomes invalid only when both counts exceed k.

For example, if k = 2, then:

• "00011" is valid because it contains only two 1s.
• "001111" is valid because it contains only two 0s.
• "000111" is invalid because it contains three 0s and three 1s.

The input size is large:

• |s| ≤ 100,000
• |queries| ≤ 100,000

A naive solution that examines every substring for every query is completely infeasible. Even a single query can contain roughly O(n²) substrings.

The challenge is therefore to preprocess the string so that each query can be answered efficiently.

Important edge cases include:

• k may be as large as n, making every substring valid.
• Query ranges may contain only one character.
• The string may consist entirely of zeros or entirely of ones.
• Many queries may overlap heavily.
• A substring becomes invalid only when both character counts exceed k, which is easy to misinterpret.

Approaches

Brute Force

For a query [l, r], enumerate every substring inside that interval.

For each substring, count zeros and ones and check whether the k-constraint holds.

This is correct because every candidate substring is explicitly examined.

However, a range of length m contains m(m+1)/2 substrings. With up to 100,000 queries, this approach is far too slow.

Even after using prefix sums to compute character counts in O(1), we still need to inspect all substrings.

Key Observation

A substring is invalid only if:

• zeros > k
• ones > k

Using a sliding window, we can determine for every starting position i the furthest ending position that still produces a valid substring.

Let:

• rightBound[i] = first position where validity breaks
• equivalently, every ending position < rightBound[i] forms a valid substring starting at i

This transforms the problem into a geometric counting problem.

For a query [l, r], we need to count pairs:

(i, j)
l ≤ i ≤ j ≤ r
j < rightBound[i]

If we define:

limit[i] = min(rightBound[i] - 1, r)

then the contribution of start index i is:

limit[i] - i + 1

The difficulty is answering this efficiently for many queries.

The crucial insight is that:

• rightBound[i] is nondecreasing.
• We can preprocess prefix sums of contributions.
• For each query, binary search separates indices whose valid region fully covers r from those clipped by r.

This yields an O((n + q) log n) solution.

Approach Comparison

Approach	Time Complexity	Space Complexity	Notes
Brute Force	O(q · n²)	O(1)	Enumerates every substring for every query
Optimal	O((n + q) log n)	O(n)	Sliding window, prefix sums, binary search

Algorithm Walkthrough

Step 1: Compute the first invalid ending for every start position

Use a sliding window.

Maintain:

• left pointer
• right pointer
• count of zeros
• count of ones

Expand the window to the right.

Whenever both counts exceed k, the window becomes invalid. At that moment:

• record rightBound[left] = right
• move left forward until validity is restored

After processing the entire string, any remaining start positions never encountered an invalid window, so:

rightBound[i] = n

for those indices.

Thus every substring:

[i, j]

is valid exactly when:

j < rightBound[i]

Step 2: Build contribution values

For each index:

fullCount[i] = rightBound[i] - i

This equals the number of valid substrings starting at i in the entire string.

Build prefix sums:

prefixFull

so ranges of fullCount can be summed quickly.

Step 3: Build another prefix sum of rightBound

Create:

prefixRight

where:

prefixRight[i+1] = prefixRight[i] + rightBound[i]

This will be useful when computing clipped contributions.

Step 4: Process a query

Suppose the query is:

[l, r]

We need:

Σ(min(rightBound[i]-1, r) - i + 1)

for all:

i ∈ [l, r]

Rewrite:

Σ(min(rightBound[i], r+1) - i)

Because rightBound is nondecreasing, we can binary search the first index:

p

such that:

rightBound[p] > r

Then:

• indices [l, p-1] have rightBound[i] ≤ r
• indices [p, r] have rightBound[i] > r

For the first group:

contribution = rightBound[i] - i

For the second group:

contribution = (r + 1) - i

Both sums can be computed using prefix sums.

Step 5: Compute the first part

For:

[l, p-1]

sum:

rightBound[i] - i

using:

• prefixRight
• prefix index sums

Step 6: Compute the second part

For:

[p, r]

sum:

(r + 1) - i

which becomes:

count * (r + 1) - Σi

Again, prefix sums give this in O(1).

Step 7: Add both parts

The sum of the two sections is the answer for the query.

Why it works

For every start position i, the sliding window computes the exact first ending position where the substring becomes invalid. Therefore all endings before rightBound[i] are valid and all endings at or after it are invalid.

For a query range [l, r], every valid substring corresponds uniquely to a pair (i, j) satisfying:

l ≤ i ≤ j ≤ r
j < rightBound[i]

The binary search partitions starts into those whose valid region ends before r and those whose valid region extends beyond r. The contribution formula counts exactly the number of valid endings for each start position, so summing these contributions yields the precise number of valid substrings.

Python Solution

from typing import List
from bisect import bisect_right

class Solution:
    def countKConstraintSubstrings(self, s: str, k: int, queries: List[List[int]]) -> List[int]:
        n = len(s)

        right_bound = [n] * n

        left = 0
        zeros = 0
        ones = 0

        for right, ch in enumerate(s):
            if ch == '0':
                zeros += 1
            else:
                ones += 1

            while zeros > k and ones > k:
                right_bound[left] = right

                if s[left] == '0':
                    zeros -= 1
                else:
                    ones -= 1

                left += 1

        prefix_right = [0] * (n + 1)
        prefix_idx = [0] * (n + 1)

        for i in range(n):
            prefix_right[i + 1] = prefix_right[i] + right_bound[i]
            prefix_idx[i + 1] = prefix_idx[i] + i

        answers = []

        for l, r in queries:
            p = bisect_right(right_bound, r, l, r + 1)

            part1 = (
                prefix_right[p] - prefix_right[l]
                - (prefix_idx[p] - prefix_idx[l])
            )

            cnt = r - p + 1

            part2 = (
                cnt * (r + 1)
                - (prefix_idx[r + 1] - prefix_idx[p])
            )

            answers.append(part1 + part2)

        return answers

Implementation Explanation

The sliding window computes right_bound[i], the first ending position that makes a substring starting at i invalid.

Because the left pointer only moves forward, the resulting right_bound array is naturally nondecreasing. This property enables binary search during query processing.

Two prefix arrays are constructed:

• prefix_right stores cumulative sums of right_bound.
• prefix_idx stores cumulative sums of indices.

For each query, bisect_right finds the boundary between positions whose valid region ends before or at r and those whose valid region extends beyond r.

The first group contributes right_bound[i] - i.

The second group contributes (r + 1) - i.

Both sums are obtained from the prefix arrays in constant time after the binary search.

Go Solution

package main

import "sort"

func countKConstraintSubstrings(s string, k int, queries [][]int) []int64 {
	n := len(s)

	rightBound := make([]int, n)
	for i := range rightBound {
		rightBound[i] = n
	}

	left := 0
	zeros := 0
	ones := 0

	for right := 0; right < n; right++ {
		if s[right] == '0' {
			zeros++
		} else {
			ones++
		}

		for zeros > k && ones > k {
			rightBound[left] = right

			if s[left] == '0' {
				zeros--
			} else {
				ones--
			}

			left++
		}
	}

	prefixRight := make([]int64, n+1)
	prefixIdx := make([]int64, n+1)

	for i := 0; i < n; i++ {
		prefixRight[i+1] = prefixRight[i] + int64(rightBound[i])
		prefixIdx[i+1] = prefixIdx[i] + int64(i)
	}

	ans := make([]int64, len(queries))

	for qi, q := range queries {
		l, r := q[0], q[1]

		p := sort.Search(r-l+1, func(x int) bool {
			return rightBound[l+x] > r
		}) + l

		part1 := (prefixRight[p] - prefixRight[l]) -
			(prefixIdx[p] - prefixIdx[l])

		cnt := int64(r - p + 1)

		part2 := cnt*int64(r+1) -
			(prefixIdx[r+1] - prefixIdx[p])

		ans[qi] = part1 + part2
	}

	return ans
}

Go-Specific Notes

The answer can be as large as approximately n²/2, which exceeds the range of a 32-bit integer. Therefore all prefix sums and final answers use int64.

Go does not provide a direct equivalent of Python's bisect_right with bounds, so sort.Search is used to locate the first index whose rightBound value exceeds r.

Worked Examples

Example 1

Input:

s = "0001111"
k = 2
queries = [[0,6]]

Sliding Window Result

i	rightBound[i]
0	5
1	6
2	7
3	7
4	7
5	7
6	7

Valid counts starting at each position:

i	rightBound[i] - i
0	5
1	5
2	5
3	4
4	3
5	2
6	1

Total:

5 + 5 + 5 + 4 + 3 + 2 + 1 = 25

For query [0,6]:

p = first index with rightBound > 6 = 2

First section:

(5-0) + (6-1) = 10

Second section:

(7-2) + (7-3) + (7-4) + (7-5) + (7-6)
= 5 + 4 + 3 + 2 + 1
= 15

Answer:

10 + 15 = 25

Example 2

Input:

s = "010101"
k = 1

Computed array:

i	rightBound[i]
0	3
1	4
2	5
3	6
4	6
5	6

Query [0,5]

Binary search:

p = 3

First part:

(3-0)+(4-1)+(5-2)=9

Second part:

(6-3)+(6-4)+(6-5)=6

Answer:

15

Query [1,4]

Every valid substring has length at most 3.

Counting contributions:

3 + 3 + 2 + 1 = 9

Answer:

9

Query [2,3]

Substrings:

"0"
"1"
"01"

All valid.

Answer:

3

Complexity Analysis

Measure	Complexity	Explanation
Time	O(n + q log n)	Sliding window preprocessing is linear, each query performs one binary search
Space	O(n)	Stores rightBound and prefix sums

The sliding window advances each pointer at most n times, giving O(n) preprocessing. Each query requires only a binary search on the monotonic rightBound array plus constant-time prefix-sum calculations, resulting in O(log n) per query.

Test Cases

from typing import List

sol = Solution()

assert sol.countKConstraintSubstrings(
    "0001111", 2, [[0, 6]]
) == [25]  # example 1

assert sol.countKConstraintSubstrings(
    "010101", 1, [[0, 5], [1, 4], [2, 3]]
) == [15, 9, 3]  # example 2

assert sol.countKConstraintSubstrings(
    "0", 1, [[0, 0]]
) == [1]  # single character

assert sol.countKConstraintSubstrings(
    "11111", 1, [[0, 4]]
) == [15]  # all substrings valid

assert sol.countKConstraintSubstrings(
    "00000", 2, [[1, 3]]
) == [6]  # all substrings inside subrange valid

assert sol.countKConstraintSubstrings(
    "01", 1, [[0, 1]]
) == [3]  # smallest mixed string

assert sol.countKConstraintSubstrings(
    "0011", 1, [[0, 3]]
) == [9]  # some invalid substrings appear

assert sol.countKConstraintSubstrings(
    "01010", 5, [[0, 4]]
) == [15]  # k larger than counts

assert sol.countKConstraintSubstrings(
    "0101010101", 1, [[0, 9], [3, 7]]
) == [27, 12]  # larger alternating pattern

assert sol.countKConstraintSubstrings(
    "101010", 1, [[0, 0], [5, 5]]
) == [1, 1]  # boundary queries

Test Summary

Test	Why
"0001111"	Validates sample behavior with only two invalid substrings
"010101"	Validates alternating pattern and multiple queries
"0"	Minimum possible string
"11111"	Single-character type, every substring valid
"00000" subrange	Query entirely inside larger string
"01"	Smallest mixed binary string
"0011"	Contains invalid longer substrings
Large k	Every substring should be valid
Long alternating string	Stress test for binary search partition
Single-element queries	Validates range boundaries

Edge Cases

Large k Makes Everything Valid

If k is at least the total number of zeros or ones in every substring, the constraint is always satisfied. In this situation the sliding window never encounters an invalid state, so every rightBound[i] remains n. The query formula naturally reduces to counting all substrings inside the range.

Single Character Query Ranges

A query such as [i, i] contains exactly one substring. Off-by-one mistakes commonly appear here. The implementation handles this correctly because the contribution formula becomes exactly one valid ending position.

Strings Containing Only One Character Type

For strings like "000000" or "111111", one character count is always zero. Since zero is always at most k, every substring is valid. The sliding window never shrinks, rightBound[i] remains n, and all substrings are counted correctly.

Boundary Where Validity First Breaks

The most subtle case occurs when a window transitions from valid to invalid because both counts become greater than k. The algorithm stores the first invalid position in rightBound[i], not the last valid position. This convention makes the validity condition simply:

j < rightBound[i]

which avoids off-by-one errors and enables the clean prefix-sum formulation used during query processing.