Learn AI Create the Future — with Pranay
Main site

LeetCode solutions

3546. Count Substrings That Satisfy K-Constraint II

My accepted Python solution to LeetCode problem 3546, Count Substrings That Satisfy K-Constraint II, running in 591ms.

  • Difficulty: Hard
  • Python
  • Runtime 591ms
  • Memory 61.3MB

Read the problem on LeetCode View on GitHub

The problem statement is LeetCode’s and stays on their site. What follows is my accepted solution.

Python

Accepted on LeetCode — runtime 591ms, memory 61.3MB, accepted 2025-12-29.

python
class Solution:
    def countKConstraintSubstrings(self, s: str, k: int, queries: List[List[int]]) -> List[int]:
        n = len(s)
        
        # For each starting position i, find the max ending position j such that s[i:j+1] satisfies k-constraint
        
        right = [0] * n
        zeros = 0
        ones = 0
        j = 0
        for i in range(n):
            while j < n and (zeros + (s[j] == '0') <= k or ones + (s[j] == '1') <= k):
                if s[j] == '0':
                    zeros += 1
                else:
                    ones += 1
                j += 1
            right[i] = j - 1
            
            if s[i] == '0':
                zeros -= 1
            else:
                ones -= 1
        
        # Precompute prefix[i] = sum of (right[j] - j + 1) for j in [0, i-1]
        prefix = [0] * (n + 1)
        for i in range(n):
            prefix[i + 1] = prefix[i] + (right[i] - i + 1)
        
        result = []
        for l, r in queries:
            # For starting positions i in [l, r]:
            # - If right[i] >= r: can end anywhere in [i, r], contributing (r - i + 1) substrings
            # - If right[i] < r: can end anywhere in [i, right[i]], contributing (right[i] - i + 1) substrings
            
            # Find the first i in [l, r] where right[i] < r (i.e., right[i] <= r - 1)
            # Note: right is monotonically non-decreasing (since we use sliding window)
            
            # Binary search for first position p in [l, r] where right[p] >= r
            # Wait no... right[i] is the max endpoint. If right[i] >= r, we use all of [i, r]
            # If right[i] < r, we use only [i, right[i]]
            
            # Since right is non-decreasing, once right[i] >= r, all subsequent right[j] >= r for j > i
            # So we find first p where right[p] >= r
            # For i in [l, p-1]: right[i] < r, use prefix sum
            # For i in [p, r]: right[i] >= r, use arithmetic sum
            
            # Actually I had it backwards. Let me fix:
            # Binary search for first p in [l, r+1] where right[p] >= r
            lo, hi = l, r + 1
            while lo < hi:
                mid = (lo + hi) // 2
                if right[mid] >= r:
                    hi = mid
                else:
                    lo = mid + 1
            p = lo  # First position with right[p] >= r
            
            # Part 1: positions l to p-1, where right[i] < r
            # count = sum of (right[i] - i + 1) for i in [l, p-1]
            if p > l:
                part1 = prefix[p] - prefix[l]
            else:
                part1 = 0
            
            # Part 2: positions p to r, where right[i] >= r
            # count = sum from i=p to r of (r - i + 1)
            num_full = r - p + 1
            if num_full > 0:
                # Arithmetic series: (r-p+1) + (r-p) + ... + 1
                part2 = num_full * (num_full + 1) // 2
            else:
                part2 = 0
            
            result.append(part1 + part2)
        
        return result

Source