Number of substrings of binary string containing K ones

In this post we will write a program in C++ to count the number of substrings of a binary string containing K ones, first we will discuss the problem with examples and then implement the solution in C++.

Given: A binary string of length N and an integer K.

Problem: Find out how many substrings of the given binary string exists which contains exactly K ones.

Example:

Input : s = “10010”
        K = 1
Output : 9
The 9 substrings containing one 1 are,
“1”, “10”, “100”, “001”, “01”, “1”, 
“10”, “0010” and “010”

Solution:

In this problem we need to find count of substrings which contains exactly K ones or in other words sum of digits in those substring is K. We first create a prefix sum array and loop over that and stop when sum value is greater than or equal to K. Now if sum at current index is (K + a) then we know that substring sum, from all those indices where sum is (a), till current index will be K, so count of indices having sum (a), will be added to result. This procedure is explained with an example below:

string s = “100101”
K = 2
prefix sum array = [1, 1, 1, 2, 2, 3]

So, at index 3,    we have prefix sum 2, 
Now total indices from where sum is 2, is 1
so result = 1

Substring considered = [“1001”]
At index 4,        we have prefix sum 2,
Now total indices from where sum is 2, is 
1 so result = 2

Substring considered = [“1001”, “10010”]
At index 5,     we have prefix sum 3,
Now total indices from where sum is 2, 
is 3 so result = 5
Substring considered = [“1001”, “10010”, 
“00101”, “0101”, “101”]

So we need to track two things, prefix sum and frequency of particular sum. In below code, instead of storing complete prefix sum, only prefix sum at current index is stored using one variable and frequency of sums in stored in an array.

Program to count the number of substrings of binary string containing K ones in C++:

You can also execute this code on our online compiler.

#include <bits/stdc++.h> 
using namespace std; 

// method returns total number of substring having K ones 
int countOfSubstringWithKOnes(string s, int K) 
{ 
    int N = s.length(); 
    int res = 0; 
    int countOfOne = 0; 
    int freq[N + 1] = {0}; 

    // initialize index having zero sum as 1 
    freq[0] = 1; 

    // loop over binary characters of string 
    for (int i = 0; i < N; i++) { 

        // update countOfOne variable with value of ith character 
        countOfOne += (s[i] - '0'); 

        // if value reaches more than K, then update result 
        if (countOfOne >= K) { 

            // add frequency of indices, having sum (current sum - K), to the result 
            res += freq[countOfOne - K]; 
        } 

        // update freqency of one's count 
        freq[countOfOne]++; 
    } 
    return res; 
} 

int main() 
{ 
    string s = "10010"; 
    int K = 1; 
    cout << countOfSubstringWithKOnes(s, K) << endl; 
    return 0; 
} 

Output:

9

Time Complexity: O(N)

Find more string related problems: String Problems

Comment below in case of any suggestion or you want to discuss more on this topic.