5643. 将数组分成三个子数组的方案数

5643. 将数组分成三个子数组的方案数

两个都是因为边界条件，wa掉了全零的情况，比如[0,0,0,0]

两次二分

class Solution {
public:
    int waysToSplit(vector<int>& nums) {
        vector<int>sum;
        sum.push_back(0);
        for(auto i:nums)sum.push_back(sum.back()+i);
        const int MOD=1e9+7;
        int n=sum.size();
        int res=0;
        for(int i=1;i<n-2;i++){
            int left=sum[i];
            int l=i+1,r=n-2;
            while(l<r){
                int m=(l+r)>>1;
                int mid=sum[m]-sum[i];
                if(mid<left)l=m+1;
                else r=m;
            }
            int mid=sum[l]-sum[i];
            int right=sum[n-1]-sum[l];
            if(left>mid || mid>right)continue;
            int e=l;

            if(mid>=left){
                l=e,r=n-2;
                while(l<r){
                    int m=(l+r+1)>>1;
                    int mid=sum[m]-sum[i];
                    int right=sum[n-1]-sum[m];
                    if(mid<=right)l=m;
                    else r=m-1;
                }
                int mid=sum[l]-sum[i];
                int right=sum[n-1]-sum[l];
                if(left>mid || mid>right)continue;
                res=(res+l-e+1)%MOD;
            }
        }
        return res;
    }
};

双指针

class Solution {
public:
    int waysToSplit(vector<int>& nums) {
        int n=nums.size();
        int mod=1e9+7;
        int sum[n+1];
        memset(sum,0,sizeof sum);
        int res=0;
        for(int i=0;i<n;i++)sum[i+1]=sum[i]+nums[i];
        for(int i=2,j=1,k=1;i<n;i++){
            while(sum[i]-sum[j]>sum[n]-sum[i])j++;
            while(k<i && sum[k]<=sum[i]-sum[k])k++;
            int lj=sum[j],mj=sum[i]-sum[j],rj=sum[n]-sum[i];
            k--;
            k=max(j,k);
            int lk=sum[k],mk=sum[i]-sum[k],rk=sum[n]-sum[i];
            if(lj<=mj && mj<=rj && lk<=mk && mk<=rk){
                res=(res+k-j+1)%mod;
            }
        }
        return res;
    }
};