<p>You are given a binary string <code>s</code>.</p>

<p>Return the number of <span data-keyword="substring-nonempty">substrings</span> with <strong>dominant</strong> ones.</p>

<p>A string has <strong>dominant</strong> ones if the number of ones in the string is <strong>greater than or equal to</strong> the <strong>square</strong> of the number of zeros in the string.</p>

<p>&nbsp;</p>
<p><strong class="example">Example 1:</strong></p>

<div class="example-block">
<p><strong>Input:</strong> <span class="example-io">s = &quot;00011&quot;</span></p>

<p><strong>Output:</strong> <span class="example-io">5</span></p>

<p><strong>Explanation:</strong></p>

<p>The substrings with dominant ones are shown in the table below.</p>
</div>

<table>
	<thead>
		<tr>
			<th>i</th>
			<th>j</th>
			<th>s[i..j]</th>
			<th>Number of Zeros</th>
			<th>Number of Ones</th>
		</tr>
	</thead>
	<tbody>
		<tr>
			<td>3</td>
			<td>3</td>
			<td>1</td>
			<td>0</td>
			<td>1</td>
		</tr>
		<tr>
			<td>4</td>
			<td>4</td>
			<td>1</td>
			<td>0</td>
			<td>1</td>
		</tr>
		<tr>
			<td>2</td>
			<td>3</td>
			<td>01</td>
			<td>1</td>
			<td>1</td>
		</tr>
		<tr>
			<td>3</td>
			<td>4</td>
			<td>11</td>
			<td>0</td>
			<td>2</td>
		</tr>
		<tr>
			<td>2</td>
			<td>4</td>
			<td>011</td>
			<td>1</td>
			<td>2</td>
		</tr>
	</tbody>
</table>

<p><strong class="example">Example 2:</strong></p>

<div class="example-block">
<p><strong>Input:</strong> <span class="example-io">s = &quot;101101&quot;</span></p>

<p><strong>Output:</strong> <span class="example-io">16</span></p>

<p><strong>Explanation:</strong></p>

<p>The substrings with <strong>non-dominant</strong> ones are shown in the table below.</p>

<p>Since there are 21 substrings total and 5 of them have non-dominant ones, it follows that there are 16 substrings with dominant ones.</p>
</div>

<table>
	<thead>
		<tr>
			<th>i</th>
			<th>j</th>
			<th>s[i..j]</th>
			<th>Number of Zeros</th>
			<th>Number of Ones</th>
		</tr>
	</thead>
	<tbody>
		<tr>
			<td>1</td>
			<td>1</td>
			<td>0</td>
			<td>1</td>
			<td>0</td>
		</tr>
		<tr>
			<td>4</td>
			<td>4</td>
			<td>0</td>
			<td>1</td>
			<td>0</td>
		</tr>
		<tr>
			<td>1</td>
			<td>4</td>
			<td>0110</td>
			<td>2</td>
			<td>2</td>
		</tr>
		<tr>
			<td>0</td>
			<td>4</td>
			<td>10110</td>
			<td>2</td>
			<td>3</td>
		</tr>
		<tr>
			<td>1</td>
			<td>5</td>
			<td>01101</td>
			<td>2</td>
			<td>3</td>
		</tr>
	</tbody>
</table>

<p>&nbsp;</p>
<p><strong>Constraints:</strong></p>

<ul>
	<li><code>1 &lt;= s.length &lt;= 4 * 10<sup>4</sup></code></li>
	<li><code>s</code> consists only of characters <code>&#39;0&#39;</code> and <code>&#39;1&#39;</code>.</li>
</ul>

class Solution:
    # exploit the fact that number of zeros is square root.
    # we can fix number of zeros. (at most sqrt(n)) = K
    # now, easy to sliding window
    # we just want the number of substrings with num of 1s >= K * K
    # sliding window: while substring num of zeros > K... increase left
    # otherwise if substring num of zeros == K, increment if 
    
    # actually sliding window is hard
    # think of starting at every index i
    # we can binary search for the range of substrings with equal to K zeros.
    # and then binary search for in that range, how many have ones >= K*K
    # 11101101110
    # 
    # the idea for sliding window is like:
    # WE WANT TO COUNT ALL VALID ANSWERS THAT END AT INDEX R...
    # we have some excess amount, however we can't necessarily use that all because if we use too much excess, we could have less than K zeros now in that substring.
    # so we count the prefixes, until we find a 0
    # 
    def numberOfSubstrings(self, s: str) -> int:
        from bisect import bisect_left, bisect_right
        from math import sqrt
        
        n = len(s)
        
        # Prefix arrays
        prefix_zeros = [0] * (n + 1)
        prefix_ones = [0] * (n + 1)
        
        for i in range(n):
            prefix_zeros[i] = prefix_zeros[i-1] + int(s[i] == '0')
            prefix_ones[i] = prefix_ones[i-1] + int(s[i] == '1')
            # prefix_zeros[i + 1] = prefix_zeros[i] + (1 if s[i] == '0' else 0)
            #prefix_ones[i + 1] = prefix_ones[i] + (1 if s[i] == '1' else 0)
        
        def count_zeros(l, r):
            return prefix_zeros[r + 1] - prefix_zeros[l]
        
        def count_ones(l, r):
            return prefix_ones[r + 1] - prefix_ones[l]
        
        # total_count = 0
        
        # Iterate over possible number of zeros K
        res = 0
        for K in range(int(sqrt(n)) + 1):
            required_ones = K * K
            # zero_positions = []
            
            # Collect positions where zeros occur
            # for i in range(n):
            #     if s[i] == '0':
            #         zero_positions.append(i)
            
            # For each valid start index
            for start in range(n):
                l = start
                r = n-1
                while l < r:
                    m = l + (r-l)//2
                    if prefix_zeros[m] - prefix_zeros[start-1] >= K:
                        r = m
                    else:
                        l = m + 1
                
                # l is the answer
                L = l
                
                #print(K, start, L, prefix_zeros[L] - prefix_zeros[start-1], prefix_zeros)
                if prefix_zeros[L] - prefix_zeros[start-1] != K: continue
                #print(K, start, L, prefix_zeros[L] - prefix_zeros[start-1])
                
                l = start
                r = n-1
                while l < r:
                    m = ceil(l + (r-l)/2)
                    if prefix_zeros[m] - prefix_zeros[start-1] <= K:
                        l = m
                    else:
                        r = m - 1
                # r is the answer
                R = r
                
                # start to ending index in [L,R] has == K*K zeros.
                # UNLESS its impossible
                
                l = L
                r = R
                while l < r:
                    m = l + (r-l)//2
                    if prefix_ones[m] - prefix_ones[start-1] >= K*K:
                        r = m
                    else:
                        l = m + 1
                # all substrings ending from [l, R] works.
                # BUT IT ALSO POSSIBLE FOR NONE TO WORK, IF THERE IS NOT ENOUGH 1'S!!!!
                
                if prefix_ones[l] - prefix_ones[start-1] >= K*K:
                    #print(K, start, l,R, R-l+1, res)
                    if l <= R:
                        res += R-l+1
            
                
                
                    
                
        return res
    
    def numberOfSubstrings(self, s: str) -> int:
        n = len(s)
        
        prefix_zeros = [0] * (n + 1)
        prefix_ones = [0] * (n + 1)
        
        for i in range(n):
            prefix_zeros[i] = prefix_zeros[i-1] + int(s[i] == '0')
            prefix_ones[i] = prefix_ones[i-1] + int(s[i] == '1')
        
        zeros = []
        for i in range(n):
            if s[i] == '0':
                zeros.append(i)
                
        res = 0
        for K in range(int(sqrt(n))+1):
            num_zeros = 0
            l = 0
            idx_zeros = 0
            #prefix_ones = 0 # BEFORE THE first 0 that is after or equal to index l
            for r in range(n):
                num_zeros += s[r] == '0'
                
                while l < r and num_zeros > K:
                    
                    num_zeros -= s[l] == '0'
                    
                    if s[l] == '0':
                        idx_zeros += 1
                    l += 1
                
                # imagine prefix_ones[r] - prefix_ones[l-1] = 12
                # K*K = 9
                # ones = 1
                
                if l <= r and num_zeros == K and prefix_ones[r] - prefix_ones[l-1] >= K*K:
                    ones = (prefix_ones[zeros[idx_zeros]] if idx_zeros < len(zeros) else prefix_ones[zeros[-1]] if zeros else 0) - (prefix_ones[zeros[idx_zeros-1]] if idx_zeros-1 >= 0 else 0)
                    #print(K,l,r,idx_zeros, zeros, ones)
                    #print(K, l, r, max(prefix_ones[r] - prefix_ones[l-1] - K*K, prefix_ones[r] - prefix_ones[l-1] - ones))
                    
                    excess = prefix_ones[r] - prefix_ones[l-1] - K*K #+ (1 if K > 0 else 0) 
                    #print(ones, excess)
                    
                    # sooo many edge cases...
                    if K == 0:
                        res += excess
                    else:
                        
                        # if ones == 0:
                        #     res += 1
                        # else:
                        res += min(ones, excess) + 1
                    #res += #excess if excess <= ones else ones
                    
                    # if ones < prefix_ones[r] - prefix_ones[l-1] - K*K:
                    #     res += ones + 1
                    # else:
                    #     res += prefix_ones[r] - prefix_ones[l-1] - K*K + 1
                    #print(res)
#                     res += max(prefix_ones[r] - prefix_ones[l-1] - K*K , prefix_ones[r] - prefix_ones[l-1] - K*K - ones)
                    
                    
                    
        return res
                
                
                
        
        
        
        
        
        


20800