LeetCode in Kotlin

3354. Make Array Elements Equal to Zero

Easy

You are given an integer array nums.

Start by selecting a starting position curr such that nums[curr] == 0, and choose a movement direction of either left or right.

After that, you repeat the following process:

A selection of the initial position curr and movement direction is considered valid if every element in nums becomes 0 by the end of the process.

Return the number of possible valid selections.

Example 1:

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

Output: 2

Explanation:

The only possible valid selections are the following:

Example 2:

Input: nums = [2,3,4,0,4,1,0]

Output: 0

Explanation:

There are no possible valid selections.

Constraints:

Solution

import kotlin.math.abs

class Solution {
    fun countValidSelections(nums: IntArray): Int {
        val rightSum = IntArray(nums.size)
        val leftSum = IntArray(nums.size)
        var result = 0
        leftSum[0] = 0
        rightSum[nums.size - 1] = 0
        for (i in 1..<nums.size) {
            leftSum[i] = leftSum[i - 1] + nums[i - 1]
        }
        for (j in nums.size - 2 downTo 0) {
            rightSum[j] = rightSum[j + 1] + nums[j + 1]
        }
        for (k in nums.indices) {
            if (nums[k] == 0 && abs((rightSum[k] - leftSum[k])) == 1) {
                result++
            }
            if (nums[k] == 0 && abs((rightSum[k] - leftSum[k])) == 0) {
                result += 2
            }
        }
        return result
    }
}