930. 和相同的二元子数组

Posted 2021-08-06 ai52learn

给你一个二元数组 nums ，和一个整数 goal ，请你统计并返回有多少个和为 goal 的 非空 子数组。

子数组 是数组的一段连续部分。

示例 1：

输入：nums = [1,0,1,0,1], goal = 2
输出：4
解释：
有 4 个满足题目要求的子数组：[1,0,1]、[1,0,1,0]、[0,1,0,1]、[1,0,1]
示例 2：

输入：nums = [0,0,0,0,0], goal = 0
输出：15
 

提示：

1 <= nums.length <= 3 * 104
nums[i] 不是 0 就是 1
0 <= goal <= nums.length

class Solution {
public:
    int numSubarraysWithSum(vector<int>& nums, int goal) {
        int l = 0, r = 0,cnt=nums[0],ret=0,a[30001]={0},b[30001]={0};
        for(int i =0;i<nums.size();i++)
        {
            int j = i;
            while(j<nums.size()&&nums[j]==0)
            {
                a[i]++;
                j++;
            }
            if(a[i]>0)
            {
                b[j-1]=a[i];
                i = j;
            }
        }
        if(!goal)
        {
            for(int i =0;i<nums.size();i++)
                if(a[i]) ret+=a[i]*(a[i]-1)/2+a[i];
            return ret;
        }
        while(l<nums.size()&&nums[l]==0) l++;
        r = l;
        if(l<nums.size()) cnt = nums[l];
        int gg = 0;
        while(l<nums.size())
        {
            while(r<nums.size()&&cnt<goal)
            {
                cnt+=nums[++r];
            }
            if(r>=nums.size()) break;
            if(goal==cnt)
            {
                int t = 1;
                if(l-1>=0&&b[l-1]>0) t*=(b[l-1]+1);
                if(r+1<=nums.size()&&a[r+1]>0) t*=(a[r+1]+1);
                ret+=t;
            }
            cnt -= nums[l++];
            while(l<nums.size()&&nums[l]==0) l++;
        }
        return ret;
    }
};