LeetCode in Kotlin

3728. Stable Subarrays With Equal Boundary and Interior Sum

Medium

You are given an integer array capacity.

A non-empty subarrays capacity[l..r] is considered stable if:

Return an integer denoting the number of stable subarrays.

Example 1:

Input: capacity = [9,3,3,3,9]

Output: 2

Explanation:

Example 2:

Input: capacity = [1,2,3,4,5]

Output: 0

Explanation:

No subarray of length at least 3 has equal first and last elements, so the answer is 0.

Example 3:

Input: capacity = [-4,4,0,0,-8,-4]

Output: 1

Explanation:

[-4,4,0,0,-8,-4] is stable because the first and last elements are both -4, and the sum of the elements strictly between them is 4 + 0 + 0 + (-8) = -4

Constraints:

Solution

class Solution {
    fun countStableSubarrays(capacity: IntArray): Long {
        var sum: Long = 0
        val map: MutableMap<Int, MutableMap<Long, Int>> = HashMap()
        var index = 0
        var ans: Long = 0
        for (c in capacity) {
            sum += c.toLong()
            var elementMap = map[c]
            if (elementMap == null) {
                elementMap = HashMap()
                map[c] = elementMap
                elementMap[sum] = 1
            } else {
                var orDefault = elementMap.getOrDefault(sum - 2 * c, 0)
                elementMap[sum] = elementMap.getOrDefault(sum, 0) + 1
                if (c == 0 && capacity[index - 1] == 0) {
                    orDefault--
                }
                ans += orDefault.toLong()
            }
            index++
        }
        return ans
    }
}