3234. 统计 1 显著的字符串的数量

题目描述

给你一个二进制字符串 s。

请你统计并返回其中 1 显著 的子字符串的数量。

如果字符串中 1 的数量 大于或等于 0 的数量的 平方，则认为该字符串是一个 1 显著 的字符串 。

思路

一个很显然的思路是，我们要枚举起点$l$，找到所有满足条件的$r$，如果暴力枚举，时间复杂度是 $O(n^2)$，但是我们在枚举r的过程中，如果目前统计的0的数量的平方已经超过所有1的数量，那后面的r肯定是不满足条件的，就不需要考虑，所以复杂度应该是 $O(nsqrt(n))$

写起来极其麻烦

class Solution {
public:
    int numberOfSubstrings(string s) {
        int n = s.size();
        s = ' ' + s;
        vector<int>pre0(n + 2), pre1(n + 2);
        vector<int>pos;
        for(int i = 1; i <= n; ++i){
            pre0[i] = pre0[i - 1] + (s[i] == '0' ? 1 : 0);
            pre1[i] = pre1[i - 1] + (s[i] == '1' ? 1 : 0);
            if(s[i] == '0')pos.push_back(i);
        }
        pre0[n + 1] = pre0[n];
        pre1[n + 1] = pre1[n];
        pos.push_back(n + 1);
        int ans = 0;
        for(int i = 1; i <= n; ++i){//枚举起点
            int id = lower_bound(pos.begin(), pos.end(), i + 1) - pos.begin();//找到下一个0的位置
            int pre_id = i - 1;
            int num0 = s[i] == '0';
            for(int j = id; j < pos.size(); ++j){
                int k = pos[j];
                if(num0 == 0){
                    ans += max(0, pre1[k] - pre1[pre_id]);
                }
                else{
                    ans +=  min(k - pre_id,  max(0, pre1[k] - pre1[i - 1] - num0 * num0 + 1));
                } 
                ++num0;
                if(num0 * num0 > pre1[n])break;
                pre_id = k;
            }
        }
        return ans;
    }
};