LeetCode 将数组分成三个子数组的方案数

class Solution {
    
    
public:
    int waysToSplit(vector<int>& nums) {
    
    
        
        int n = nums.size();
        vector<int> presum(n + 1, 0);

        int MOD = 1e9 + 7;

        // 计算前缀和
        for (int i = 1; i <= n; i++)
        {
    
    
            presum[i] = presum[i - 1] + nums[i - 1];
        }

        int ret = 0;

        // i < n - 2 是因为后两个数组要各至少有一个元素
        for (int i = 1; i < n; i++)
        {
    
    
            // 不理解
            if (presum[i] * 3 > presum[n]) break;

            // 找到第一个使得 第二个子数组的元素和 大于等于 第一个子数组的元素和 的元素,
            // 
            int l = i + 1, r = n - 1;
            while (l <= r)
            {
    
    
                int mid = (l + r) / 2;
                // 如果第二个子数组的元素和 大于等于 第一个子数组的元素和
                // 说明当前mid满足条件,第二个子数组可以减少元素
                if (presum[mid] - presum[i] >= presum[i])
                {
    
    
                    r = mid - 1;
                }
                // 如果第二个子数组的元素和 小于 第一个子数组的元素和
                // 说明当前mid不满足条件,需要为第二个子数组增加元素
                else
                {
    
    
                    l = mid + 1;
                }
            }

            //
            int ll = i + 1, rr = n - 1;
            while (ll <= rr)
            {
    
    
                int mid = (ll + rr) / 2;
                // 如果第三个子数组的元素和 大于等于 第二个子数组的元素和
                if (presum[n] - presum[mid] >= presum[mid] - presum[i])
                {
    
    
                    ll = mid + 1;
                }
                else
                {
    
    
                    rr = mid - 1; 
                }
            }
            ret += rr - l + 1;
 
            ret %= MOD;
        }
        return ret;
    }
};

猜你喜欢

转载自blog.csdn.net/qq_32862515/article/details/112254920
今日推荐