3234. Count the Number of Substrings With Dominant Ones

Hint

class Solution {
public:
    int numberOfSubstrings(string s)
    {
        // To store the count of valid substrings
        // Vector to store prefix sums of '1's

        // Compute prefix sums for the number of '1's
        // Initialize the first element based on the first character
        // Compute prefix sums for the rest of the array

        // Iterate over each possible starting index of the substring
        for ()
        {
            // Iterate over each possible ending index of the substring
            for ()
            {
                // Calculate the number of '1's and '0's in the current substring s[i..j]

                // CASE 1: Substring is not dominant
                if ()
                {
                    // Adjust the end pointer j to skip over non-dominant substrings
                }

                // CASE 2: Substring is exactly dominant
                else if ()
                {
                    // Count this substring as valid
                }
                // CASE 3: Substring is more dominant
                else if ()
                {
                    // Count this substring as valid
                    // Calculate how many more substrings are valid by skipping forward
                    // Determine the next position to skip to

                    // If skipping exceeds the string length, count all remaining substrings
                    if ()
                    {
                    }
                    else
                    {
                        // Otherwise, count substrings within bounds
                    }

                    // Update j to the next valid position
                }
            }
        }
        // Return the total count of valid substrings
    }
};

class Solution {
public:
    int numberOfSubstrings(string s)
    {
        int n = s.size(); // Length of the input string
        int res = 0; // To store the count of valid substrings
        vector<int> prefix(n, 0); // Vector to store prefix sums of '1's

        // Compute prefix sums for the number of '1's
        prefix[0] = (int)(s[0] - '0'); // Initialize the first element based on the first character
        for (int i = 1; i < n; i++)
        {
            prefix[i] = prefix[i - 1] + (int)(s[i] - '0'); // Compute prefix sums for the rest of the array
        }

        // Iterate over each possible starting index of the substring
        for (int i = 0; i < n; i++)
        {
            // Iterate over each possible ending index of the substring
            for (int j = i; j < n; j++)
            {
                // Calculate the number of '1's and '0's in the current substring s[i..j]
                int ones = prefix[j] - (i == 0 ? 0 : prefix[i - 1]);
                int zeros = (j - i + 1) - ones;

                // CASE 1: Substring is not dominant
                if (zeros * zeros > ones)
                {
                    // Adjust the end pointer j to skip over non-dominant substrings
                    j += (zeros * zeros - ones - 1);
                }
                // CASE 2: Substring is exactly dominant
                else if (zeros * zeros == ones)
                {
                    ++res; // Count this substring as valid
                }
                // CASE 3: Substring is more dominant
                else if (zeros * zeros < ones)
                {
                    ++res; // Count this substring as valid
                    // Calculate how many more substrings are valid by skipping forward
                    int diff = (int) sqrt(ones) - zeros;
                    int nextj = j + diff; // Determine the next position to skip to

                    // If skipping exceeds the string length, count all remaining substrings
                    if (nextj >= n)
                    {
                        res += (n - j - 1);
                    }
                    else
                    {
                        res += diff; // Otherwise, count substrings within bounds
                    }
                    // Update j to the next valid position
                    j = nextj;
                }
            }
        }
        return res; // Return the total count of valid substrings
    }
};

• T: $O(n \sqrt n)$
• S: $O(n)$