LeetCode in Kotlin

1712. Ways to Split Array Into Three Subarrays

Medium

A split of an integer array is good if:

Given nums, an array of non-negative integers, return the number of good ways to split nums. As the number may be too large, return it modulo 109 + 7.

Example 1:

Input: nums = [1,1,1]

Output: 1

Explanation: The only good way to split nums is [1] [1] [1].

Example 2:

Input: nums = [1,2,2,2,5,0]

Output: 3

Explanation: There are three good ways of splitting nums:

[1] [2] [2,2,5,0]

[1] [2,2] [2,5,0]

[1,2] [2,2] [5,0]

Example 3:

Input: nums = [3,2,1]

Output: 0

Explanation: There is no good way to split nums.

Constraints:

Solution

class Solution {
    fun waysToSplit(nums: IntArray): Int {
        var sum = 0
        for (num in nums) {
            sum += num
        }
        var cur = 0
        var res: Long = 0
        var i = 0
        var idx1 = 1
        var sum1 = nums[0]
        var idx2 = 1
        var sum2 = nums[0]
        while (i < nums.size) {
            cur += nums[i]
            val right = sum - cur
            if (i == 0 || i == nums.size - 1) {
                i++
                continue
            }
            while (idx1 <= i && sum1 <= cur - sum1) {
                sum1 += nums[idx1++]
            }
            while (idx2 < idx1 && cur - sum2 > right) {
                sum2 += nums[idx2++]
            }
            if (idx1 > idx2) {
                res = (res + idx1 - idx2) % 1000000007
            }
            i++
        }
        return res.toInt()
    }
}