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

164 阅读1分钟

1712. 将数组分成三个子数组的方案数 - 力扣(LeetCode)

我们称一个分割整数数组的方案是 好的 ,当它满足:

  • 数组被分成三个 非空  连续子数组,从左至右分别命名为  left , mid , right 。
  • left  中元素和小于等于  mid  中元素和,mid  中元素和小于等于  right  中元素和。

给你一个 非负 整数数组  nums ,请你返回  好的 分割 nums  方案数目。由于答案可能会很大,请你将结果对 109 + 7  取余后返回。

示例 1:

输入: nums = [1,1,1]
输出: 1
解释: 唯一一种好的分割方案是将 nums 分成 [1] [1] [1] 。

示例 2:

输入: nums = [1,2,2,2,5,0]
输出: 3
解释: nums 总共有 3 种好的分割方案:
[1] [2] [2,2,5,0]
[1] [2,2] [2,5,0]
[1,2] [2,2] [5,0]

示例 3:

输入: nums = [3,2,1]
输出: 0
解释: 没有好的分割方案。

提示:

  • 3 <= nums.length <= 10^5
  • 0 <= nums[i] <= 10^4

思路

本题可以用 前缀和 + 双指针 求解。设 nums 中数字和为 sum,三个非空数组和分别为 lsummsumrsum,由于 lsum <= msum <= rsum,所有 lsum<=sum / 3msum + rsum = sum - lsummsum <= (sum - lsum) / 2

我们假设 left 数组为 nums[0]nums[i]mid 数组为 nums[i+1]到 nums[j],right 数组为 nums[j + 1] 到数组末尾。用 f(i)表示 numsi + 1 个数字的和,即 f(i) = nums[0] + nums[1] + ... + nums[i]f(i) = lsum, f(j) = lsum + msum >= f(i) * 2 ,又由于 msum <= (sum - lsum ) / 2 ,故 f(j) = lsum + msum <= lsum + (sum - lsum) / 2 <= (sum + lsum) / 2 <= (sum + f(i)) / 2, 由上面的推导知 j 的范围和 f(i) 正相关,单调递增。用 startend 表示 j 的范围,本次有 end - start + 1 中分割方法。代码如下。

解题

/**
 * @param {number[]} nums
 * @return {number}
 */
var waysToSplit = function (nums) {
  const MOD = 10 ** 9 + 7;
  const n = nums.length;
  for (let i = 1; i < n; i++) {
    nums[i] += nums[i - 1];
  }

  let res = 0;
  const m = Math.floor(nums[n - 1] / 3);
  let j = 1;
  let k = 1;
  for (let i = 0; nums[i] <= m && i < n - 2; i++) {
    let left = 2 * nums[i];
    let right = (nums[n - 1] + nums[i]) / 2;
    j = Math.max(j, i + 1);
    while (j < n - 1 && nums[j] < left) {
      j++;
    }
    while (k < n - 2 && nums[k + 1] <= right) {
      k++;
    }
    res = (res + k - j + 1) % MOD;
  }
  return res;
};