Count the Number of Substrings With Dominant Ones

main.md

Question

Approach

I solve this by using the key observation that a substring with k zeros can only be dominant if k is at most √n, since the condition requires ones ≥ k² and a substring cannot have more than n ones. So I first collect the positions of all zeros, then count all-ones substrings directly (they always satisfy the condition). For each possible zero count k from 1 to √n, I slide a window of k consecutive zeros over the string. For each such window, I compute how far the substring can extend left and right using the ones around the zero segment, and I also count how many ones lie between the first and last zero in that window. This lets me determine exactly which substrings include exactly those k zeros, and how many extra ones they need to meet the condition ones ≥ k². I then count in O(1) how many left/right extensions satisfy this requirement by subtracting the failing extensions from the total possible extensions.

Implementation

class Solution:
    def numberOfSubstrings(self, s: str) -> int:
        n = len(s)
        zpos = [i for i,ch in enumerate(s) if ch == '0']
        m = len(zpos)
        
        dominant = 0
        run = 0
        for ch in s:
            if ch == '1':
                run += 1
            else:
                if run:
                    dominant += run * (run + 1) // 2
                run = 0
        if run:
            dominant += run * (run + 1) // 2
        
        if m == 0:
            return dominant
        
        K = int(math.isqrt(n))
        
        def failing_pairs(L: int, R: int, t: int) -> int:
            if t < 0:
                return 0
            total = (L + 1) * (R + 1)
            if t >= L + R:
                return total
            a = L if L < t else t
            u_val = t - R
            if u_val < 0:
                u = -1
            else:
                u = u_val if u_val <= a else a
            count1 = 0
            if u >= 0:
                count1 = (u + 1) * (R + 1)
            rem = a - u
            sum2 = 0
            if rem > 0:
                sum_x = (a * (a + 1) // 2) - (u * (u + 1) // 2)
                sum2 = rem * (t + 1) - sum_x
            return count1 + sum2
        
        for k in range(1, K + 1):
            if m < k:
                break
            for i in range(0, m - k + 1):
                left_zero_idx = i
                right_zero_idx = i + k - 1
                a = zpos[left_zero_idx]
                b = zpos[right_zero_idx]
                prev = zpos[left_zero_idx - 1] if left_zero_idx - 1 >= 0 else -1
                nxt = zpos[right_zero_idx + 1] if right_zero_idx + 1 < m else n
                L = a - prev - 1
                R = nxt - b - 1
                ones_inside = (b - a + 1) - k
                need = k * k - ones_inside
                total_pairs = (L + 1) * (R + 1)
                if need <= 0:
                    dominant += total_pairs
                else:
                    t = need - 1
                    fail = failing_pairs(L, R, t)
                    dominant += total_pairs - fail
        
        return dominant

Complexities

• Time: O(m * √n)
• Space: O(n)