Medium
You are given an integer array capacity.
A non-empty subarrays capacity[l..r] is considered stable if:
capacity[l] = capacity[r] = capacity[l + 1] + capacity[l + 2] + ... + capacity[r - 1]).Return an integer denoting the number of stable subarrays.
Example 1:
Input: capacity = [9,3,3,3,9]
Output: 2
Explanation:
[9,3,3,3,9] is stable because the first and last elements are both 9, and the sum of the elements strictly between them is 3 + 3 + 3 = 9.[3,3,3] is stable because the first and last elements are both 3, and the sum of the elements strictly between them is 3.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:
3 <= capacity.length <= 105-109 <= capacity[i] <= 109class 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
}
}