<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>

import math

class Solution:
    def numberOfSubstrings(self, s: str) -> int:
        n = len(s)
        m = ceil(math.sqrt(n))
        print(m)

        ones = 0
        zeros = 0
        cnt0 = [0]
        cnt1 = [0]

        for i in range(n):
            if s[i] == '0':
                zeros += 1
            else:
                ones += 1
            cnt0.append(zeros)
            cnt1.append(ones)

        nextZero = [0 for i in range(n+1)]
        zeroInd = -1
        nextZero[n] = -1
        for i in range(n-1, -1, -1):
            if s[i] == '0':
                zeroInd = i
            nextZero[i] = zeroInd
        # print(nextZero)
            
        answer = 0
        for l in range(n):
            prevR = l-1
            r = nextZero[l]

            cntOnesPrev = 0
            cntZerosPrev = 0

            # print("l = ", l)
            
            while r != -1 and cntZerosPrev < m:
                cntOnesNew = r - prevR - 1
                
                # Case 1. Don't include 0
                diff = cntOnesPrev + cntOnesNew - cntZerosPrev ** 2
                # print("Cnt 1: ", cntOnesPrev, cntOnesNew, cntZerosPrev, diff)
                if diff >= 0 and cntOnesNew > 0:
                    answer += min(diff + 1, cntOnesNew)
                    # print("Diff 1: ", diff, cntOnesNew)
                    # print("Case 1:", prevR, r, answer)

                cntZerosPrev += 1
                cntOnesPrev += cntOnesNew

                # Case 2. Include 0
                answer += (cntOnesPrev >= cntZerosPrev ** 2)
                # if cntOnesPrev >= cntZerosPrev ** 2:
                    # print("Case 2:", prevR, r, answer)

                prevR = r
                r = nextZero[r+1]
                
            if r == -1:
                cntOnesNew = n - prevR - 1
                diff = cntOnesPrev + cntOnesNew - cntZerosPrev ** 2
                if diff >= 0 and cntOnesNew > 0:
                    answer += min(diff + 1, cntOnesNew)
                    # print("Cnt 3: ", cntOnesPrev, cntOnesNew, cntZerosPrev, diff)
                    # print ("Case 3:", prevR, answer)
            # print()

        
        return answer
20600