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

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

• 数组被分成三个 非空 连续子数组，从左至右分别命名为 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 <= 105
• 0 <= nums[i] <= 104

题解一（前缀和 + 枚举方案 · 超出时间限制）

枚举所有方案：

class Solution {
    fun waysToSplit(nums: IntArray): Int {
        val MOD = 1000000007
        val n = nums.size
        // 前缀和
        for (i in 1 until n) {
            nums[i] += nums[i - 1]
        }
        var ret = 0L
        for (i in 1 .. n - 2) { // 在第 [i] 位前分割
            if (nums[i - 1] * 3 > nums[n - 1]) break
            for (j in i + 1 .. n - 1) { // 在第 [j] 位前分割
                val k1 = nums[i - 1]
                val k2 = nums[j - 1] - nums[i - 1]
                val k3 = nums[n - 1] - nums[j - 1]
                if (k2 in k1 .. k3) ret = (ret + 1) % MOD
            }
        }
        return ret.toInt()
    }
}

复杂度分析：

• 时间复杂度：O(n^2) 预处理前缀和时间为 O(n)，总共有 n^2 个子问题，每个子问题时间为 O(1)；
• 空间复杂度：O(1) 原地数组预处理前缀和。

题解二（前缀和 + 两次二分查找）

使用两次二分查找优化内层循环：

class Solution {
    fun waysToSplit(nums: IntArray): Int {
        val MOD = 1000000007
        val n = nums.size
        // 前缀和
        for (i in 1 until n) {
            nums[i] += nums[i - 1]
        }
        var ret = 0L
        for (i in 0 .. n - 3) { // 在第 [i] 位分割
            if (nums[i] * 3 > nums[n - 1]) break
            // x - nums[i] >= nums[i]
            // 寻找大于等于 2 * nums[i] 的第一个元素
            var left = i + 1
            var right = n - 2
            while (left < right) {
                val mid = (left + right) ushr 1
                if (nums[mid] >= 2 * nums[i]) {
                    right = mid
                } else {
                    left = mid + 1
                }
            }
            if (nums[left] < 2 * nums[i]) continue
            val from = left
            // x - nums[i] <= top - x
            // 寻找小于等于 (nums[n - 1] + nums[i]) / 2 的最后一个元素
            right = n - 2
            while (left < right) {
                val mid = (left + right + 1) ushr 1
                if (nums[mid] <= (nums[n - 1] + nums[i]) / 2) {
                    left = mid
                } else {
                    right = mid - 1
                }
            }
            if (nums[left] > (nums[n - 1] + nums[i]) / 2) continue
            val to = left
            ret = (ret + to - from + 1) % MOD
        }
        return ret.toInt()
    }
}

复杂度分析：

• 时间复杂度：O(nlgn) 预处理前缀和时间为 O(n)，内层循环时间为 O(lgn)；
• 空间复杂度：O(1) 原地数组预处理前缀和。

题解三（前缀和 + 双指针）

j 指针与 k 指针对 i 指针单调递增，注意 i、j、k 的上界。

[]

class Solution {
    fun waysToSplit(nums: IntArray): Int {
        val MOD = 1000000007
        val n = nums.size
        // 前缀和
        for (i in 1 until n) {
            nums[i] += nums[i - 1]
        }
        var ret = 0L
        var i = 0
        var j = 1
        var k = 1
        while (i < n - 2 && nums[i] * 3 <= nums[n - 1]) {
            j = max(j, i + 1)
            // 寻找大于等于 2 * nums[i] 的第一个元素
            // 寻找小于等于 (nums[n - 1] + nums[i]) / 2 的最后一个元素
            while (j < n - 1 && nums[j] < nums[i] * 2) j++
            while (k < n - 2 && nums[k + 1] - nums[i] <= nums[n - 1] - nums[k + 1]) k++
            // println("i=i, j=j, k=k")
            ret = (ret + k - j + 1) % MOD
            i++
        }
        return ret.toInt()
    }
}

[]

class Solution {

    public static final int MOD = (int) (1e9 + 7);

    public int waysToSplit(int[] nums) {
        int n = nums.length;
        // 前缀和
        int[] p = new int[n];
        for (int i = 1; i < n; i++) {
            nums[i] += nums[i - 1];
        }
        long res = 0;
        for (int i = 0, j = 1, k = 1; i < n - 2 && nums[i] * 3 <= nums[n - 1]; i++) {
            j = Math.max(i + 1, j);
            while (j < n - 1 && nums[j] < 2 * nums[i]) j++;
            while (k < n - 2 && nums[k + 1] <= (nums[n - 1] + nums[i]) / 2) k++;
            res += k - j + 1;
        }
        return (int) (res % MOD);
    }
}

复杂度分析：

• 时间复杂度：O(n) 预处理前缀和时间为 O(n)，三个指针均最多移动 n 次；
• 空间复杂度：O(1) 原地数组预处理前缀和。