Count substrings that satisfy k constraint ii

Array String Binary Search Sliding Window Prefix Sum

23.4% acceptance

Feb 25, 2026

140

You are given a binary string s, integer k, and 2D array queries = [li, ri]. Return for each query the number of substrings of s[li..ri] that satisfy the k-constraint (count_0 <= k or count_1 <= k).

Solution

Rust

Time O(n²)

Space O(n)

LeetCode

solution.rs

impl Solution {
  pub fn count_k_constraint_substrings(s: String, k: i32, queries: Vec<Vec<i32>>) -> Vec<i64> {
    let k = k as usize;
    let b: Vec<usize> = s.bytes().map(|c| (c - b'0') as usize).collect();
    let n = b.len();

    // Compute left[j]: leftmost i such that s[i..=j] satisfies k-constraint
    let mut left = vec![0usize; n];
    let mut c0 = 0usize;
    let mut c1 = 0usize;
    let mut li = 0usize;
    for j in 0..n {
      if b[j] == 0 { c0 += 1; } else { c1 += 1; }
      while c0 > k && c1 > k {
        if b[li] == 0 { c0 -= 1; } else { c1 -= 1; }
        li += 1;
      }
      left[j] = li;
    }

    // f[j] = j + 1 - left[j]: number of valid substrings ending at j (within [0..j])
    let f: Vec<i64> = (0..n).map(|j| (j + 1 - left[j]) as i64).collect();

    // Prefix sum of f
    let mut pref_f = vec![0i64; n + 1];
    for j in 0..n {
      pref_f[j + 1] = pref_f[j] + f[j];
    }

    let mut result = Vec::with_capacity(queries.len());
    for q in &queries {
      let l = q[0] as usize;
      let r = q[1] as usize;

      // Find largest p in [l, r] where left[p] <= l
      // left is non-decreasing; use binary search
      // partition_point returns first j where left[j] > l
      let cutoff = left.partition_point(|&lj| lj <= l);
      // last valid j with left[j] <= l is cutoff - 1
      let p = if cutoff == 0 {
        // No j with left[j] <= l => all use f[j]
        let sum2 = pref_f[r + 1] - pref_f[l];
        result.push(sum2);
        continue;
      } else {
        (cutoff - 1).min(r)
      };

      let sum1 = if p >= l {
        let cnt = (p - l + 1) as i64;
        cnt * (cnt + 1) / 2
      } else {
        0
      };
      let sum2_start = if p >= l { p + 1 } else { l };
      let sum2 = pref_f[r + 1] - pref_f[sum2_start];
      result.push(sum1 + sum2);
    }
    result
  }
}