930 Binary Subarrays With Sum 和相同的二元子數組
Description:
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:
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 * 10^4
nums[i] is either 0 or 1.
0 <= goal <= nums.length
題目描述:
給你一個二元數組 nums ，和一個整數 goal ，請你統計並返回有多少個和爲 goal 的 非空 子數組。
子數組 是數組的一段連續部分。
示例 :
示例 1：
輸入：nums = [1,0,1,0,1], goal = 2
輸出：4
解釋：
有 4 個滿足題目要求的子數組：[1,0,1]、[1,0,1,0]、[0,1,0,1]、[1,0,1]
示例 2：
輸入：nums = [0,0,0,0,0], goal = 0
輸出：15
提示:
1 <= nums.length <= 3 * 10^4
nums[i] 不是 0 就是 1
0 <= goal <= nums.length
思路:
前綴和 ➕ 哈希表
參考 LeetCode #560 Subarray Sum Equals K 和爲K的子數組
時間複雜度爲 O(n), 空間複雜度爲 O(n)
代碼:
C++:
class Solution
{
public:
int numSubarraysWithSum(vector<int>& nums, int goal)
{
int n = nums.size(), result = 0;
vector<int> pre(n + 1, 0);
unordered_map<int, int> m;
for (int i = 0; i < n; i++)
{
++m[pre[i]];
result += m[(pre[i + 1] = pre[i] + nums[i]) - goal];
}
return result;
}
};
Java:
class Solution {
public int numSubarraysWithSum(int[] nums, int goal) {
int n = nums.length, result = 0;
int pre[] = new int[n + 1];
Map<Integer, Integer> map = new HashMap<>();
for (int i = 0; i < n; i++) {
map.put(pre[i], map.getOrDefault(pre[i], 0) + 1);
result += map.getOrDefault((pre[i + 1] = pre[i] + nums[i]) - goal, 0);
}
return result;
}
}
Python:
class Solution:
def numSubarraysWithSum(self, nums: List[int], goal: int) -> int:
pre, result, d = [0] * ((n := len(nums)) + 1), 0, defaultdict(int)
for i in range(n):
d[pre[i]] += 1
pre[i + 1] = pre[i] + nums[i]
result += d[pre[i + 1] - goal]
return result