Medium
Given a binary array nums
and an integer goal
, return the number of non-empty subarrays with a sum goal
.
A subarray is a contiguous part of the array.
Example 1:
Input: nums = [1,0,1,0,1], goal = 2
Output: 4
Explanation: The 4 subarrays are bolded and underlined below:
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]
[1,0,1,0,1]
Example 2:
Input: nums = [0,0,0,0,0], goal = 0
Output: 15
Constraints:
1 <= nums.length <= 3 * 104
nums[i]
is either 0
or 1
.0 <= goal <= nums.length
class Solution {
fun numSubarraysWithSum(nums: IntArray, goal: Int): Int {
return atmostK(nums, goal) - atmostK(nums, goal - 1)
}
fun atmostK(arr: IntArray, k: Int): Int {
var cnt = 0
var i = 0
var j = 0
var sum = 0
while (j < arr.size) {
sum += arr[j]
if (sum > k) {
while (i <= j && sum > k) {
sum -= arr[i]
i++
}
}
cnt += j - i + 1
j++
}
return cnt
}
}