Given a string which contains only ones and zeros, return the number of substrings with ones more than zeros

Question

Suppose that S is a string containing only 0 and 1. I would like to count the number of not null substrings of S in which the number of 0s is less than the number of 1s.

Using brute-force method given below we can have an algorithm that solves this problem in O(n^2):

moreOnes(s):
count = 0
n = s.len()

for i = 1 to n
    one = 0
    zero = 0
    for j = i to n
        if s[i] == '1'
            one++
        else
            zero++
        
        if one > zero
            count++

return count

But can we have an algorithm with a better time complexity of O(n*logn) or O(n), and space complexity be anything from O(1) to O(n)?

Answer 1

Consider an array A[i] that contains the number of ones in the range 1..i minus the number of zeros in the range 1..i.

With this array it now takes only O(1) time to compute the number of ones minus the number of zeros in a substring from i..j by computing A[j]-A[i-1].

For a given endpoint j, you want to find all start points i<=j such that A[j]-A[i-1]>0. Equivalently, you want to know how many values of A[i-1] have value less than A[j].

This can be solved with Fenwick trees by:

1. Loop over j
2. Add 1 at location A[j] in Fenwick tree
3. Lookup cumulative value from Fenwick tree in range up to A[j]-1 - this is the number of substrings that end at j and satisfy the desired property of more ones than zeros.

This will take O(n) space for the Fenwick tree and O(nlogn) time as each lookup is O(logn).

(Note that A[j] may go negative, while Fenwick trees typically work with positive data. You can work around this by adding a constant offset of n to A[j] so all entries are positive.)