Count Substrings That Satisfy K-Constraint I

Problem

You are given a binary string s and an integer k.

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 denoting the number of substrings of s that satisfy the k-constraint.

Examples

Example 1

Input: s = "10101", k = 1

Output: 12

Explanation:

Every substring of `s` except the substrings `"1010"`, `"10101"`, and `"0101"`
satisfies the k-constraint.

Example 2

Input: s = "1010101", k = 2

Output: 25

Explanation:

Every substring of `s` except the substrings with a length greater than 5
satisfies the k-constraint.

Example 3

Input: s = "11111", k = 1

Output: 15

Explanation:

All substrings of `s` satisfy the k-constraint.

Constraints

• 1 <= s.length <= 50
• 1 <= k <= s.length
• s[i] is either '0' or '1'.

Solution

Method 1 – Brute Force with Prefix Sums

Intuition

For each substring, count the number of '0's and '1's. If either count is at most k, the substring satisfies the k-constraint. Since the string is short (length ≤ 50), we can check all substrings efficiently using prefix sums.

Approach

1. Precompute prefix sums for the number of '0's and '1's up to each index.
2. For every substring s[i:j], use the prefix sums to get the count of '0's and '1's in O(1) time.
3. If either count is ≤ k, increment the answer.
4. Return the total count.

Code

C++

class Solution {
public:
    int countSubstrings(string s, int k) {
        int n = s.size(), ans = 0;
        vector<int> p0(n+1), p1(n+1);
        for (int i = 0; i < n; ++i) {
            p0[i+1] = p0[i] + (s[i] == '0');
            p1[i+1] = p1[i] + (s[i] == '1');
        }
        for (int i = 0; i < n; ++i) {
            for (int j = i; j < n; ++j) {
                int c0 = p0[j+1] - p0[i];
                int c1 = p1[j+1] - p1[i];
                if (c0 <= k || c1 <= k) ans++;
            }
        }
        return ans;
    }
};

Go

type Solution struct{}
func (Solution) CountSubstrings(s string, k int) int {
    n := len(s)
    p0 := make([]int, n+1)
    p1 := make([]int, n+1)
    for i := 0; i < n; i++ {
        p0[i+1] = p0[i]
        p1[i+1] = p1[i]
        if s[i] == '0' {
            p0[i+1]++
        } else {
            p1[i+1]++
        }
    }
    ans := 0
    for i := 0; i < n; i++ {
        for j := i; j < n; j++ {
            c0 := p0[j+1] - p0[i]
            c1 := p1[j+1] - p1[i]
            if c0 <= k || c1 <= k {
                ans++
            }
        }
    }
    return ans
}

Java

class Solution {
    public int countSubstrings(String s, int k) {
        int n = s.length(), ans = 0;
        int[] p0 = new int[n+1], p1 = new int[n+1];
        for (int i = 0; i < n; ++i) {
            p0[i+1] = p0[i] + (s.charAt(i) == '0' ? 1 : 0);
            p1[i+1] = p1[i] + (s.charAt(i) == '1' ? 1 : 0);
        }
        for (int i = 0; i < n; ++i) {
            for (int j = i; j < n; ++j) {
                int c0 = p0[j+1] - p0[i];
                int c1 = p1[j+1] - p1[i];
                if (c0 <= k || c1 <= k) ans++;
            }
        }
        return ans;
    }
}

Kotlin

class Solution {
    fun countSubstrings(s: String, k: Int): Int {
        val n = s.length
        val p0 = IntArray(n+1)
        val p1 = IntArray(n+1)
        for (i in 0 until n) {
            p0[i+1] = p0[i] + if (s[i] == '0') 1 else 0
            p1[i+1] = p1[i] + if (s[i] == '1') 1 else 0
        }
        var ans = 0
        for (i in 0 until n) {
            for (j in i until n) {
                val c0 = p0[j+1] - p0[i]
                val c1 = p1[j+1] - p1[i]
                if (c0 <= k || c1 <= k) ans++
            }
        }
        return ans
    }
}

Python

class Solution:
    def countSubstrings(self, s: str, k: int) -> int:
        n = len(s)
        p0 = [0] * (n+1)
        p1 = [0] * (n+1)
        for i in range(n):
            p0[i+1] = p0[i] + (s[i] == '0')
            p1[i+1] = p1[i] + (s[i] == '1')
        ans = 0
        for i in range(n):
            for j in range(i, n):
                c0 = p0[j+1] - p0[i]
                c1 = p1[j+1] - p1[i]
                if c0 <= k or c1 <= k:
                    ans += 1
        return ans

Rust

impl Solution {
    pub fn count_substrings(s: String, k: i32) -> i32 {
        let n = s.len();
        let s = s.as_bytes();
        let mut p0 = vec![0; n+1];
        let mut p1 = vec![0; n+1];
        for i in 0..n {
            p0[i+1] = p0[i] + if s[i] == b'0' { 1 } else { 0 };
            p1[i+1] = p1[i] + if s[i] == b'1' { 1 } else { 0 };
        }
        let mut ans = 0;
        for i in 0..n {
            for j in i..n {
                let c0 = p0[j+1] - p0[i];
                let c1 = p1[j+1] - p1[i];
                if c0 <= k || c1 <= k {
                    ans += 1;
                }
            }
        }
        ans
    }
}

TypeScript

class Solution {
    countSubstrings(s: string, k: number): number {
        const n = s.length;
        const p0 = Array(n+1).fill(0);
        const p1 = Array(n+1).fill(0);
        for (let i = 0; i < n; ++i) {
            p0[i+1] = p0[i] + (s[i] === '0' ? 1 : 0);
            p1[i+1] = p1[i] + (s[i] === '1' ? 1 : 0);
        }
        let ans = 0;
        for (let i = 0; i < n; ++i) {
            for (let j = i; j < n; ++j) {
                const c0 = p0[j+1] - p0[i];
                const c1 = p1[j+1] - p1[i];
                if (c0 <= k || c1 <= k) ans++;
            }
        }
        return ans;
    }
}

Complexity

• ⏰ Time complexity: O(n^2), where n is the length of s, since we check all substrings.
• 🧺 Space complexity: O(n) for the prefix sum arrays.