LeetCode 930和相同的二元子数组 | 前缀和与哈希表引言和相同的二元子数组Binary Subarrays With Sum是 LeetCode 第 930 题难度为 Medium。题目要求在二元数组元素只有 0 和 1中找出子数组和等于 S 的数量。这道题与 LeetCode 560和为 K 的子数组非常相似但因为数组是二元的可以使用更高效的方法。对于二元数组我们可以使用计数法来优化当 S 较大时遍历所有子数组可能需要 O(n²) 时间但使用哈希表方法仍然是 O(n)。另一种方法是使用滑动窗口因为数组是二元的可以更高效地处理。问题分析题目描述给定一个二元数组 nums 和一个整数 S计算有多少个子数组的元素和等于 S。例如输入 nums [1,0,1,0,1]S 2有 4 个子数组的和等于 2[1,0,1]索引 0-2[1,0,1]索引 2-4[0,1,0,1]索引 1-4[1]但这个和为 1不是 2所以只有前三个实际上正确答案是[1,0,1]索引 0-21012[1,0,1]索引 2-41012[1,0,1,0,1]的和是3不是2问题特点因为数组是二元的元素只能是 0 或 1。这使得我们可以使用更高效的方法。但核心思路仍然是前缀和与哈希表。解决方案前缀和哈希表方法def numSubarraysWithSum(nums, s): prefix_sum 0 count 0 prefix_map {0: 1} for num in nums: prefix_sum num count prefix_map.get(prefix_sum - s, 0) prefix_map[prefix_sum] prefix_map.get(prefix_sum, 0) 1 return count这与 LeetCode 560 的解法完全相同。对于任意整数数组包括二元数组这个方法都适用。计数方法优化因为数组是二元的我们可以使用更直接的方法。观察发现对于子数组和为 S 的情况子数组中必须恰好有 S 个 1。这 S 个 1 之间的 0 可以任意数量。def numSubarraysWithSum_optimized(nums, s): def countAtMost(k): if k 0: return 0 result 0 left 0 count 0 for right in range(len(nums)): if nums[right] 1: count 1 while count k: if nums[left] 1: count - 1 left 1 result right - left 1 return result return countAtMost(s) - countAtMost(s - 1)这个方法计算至多 s 个 1的子数组数量减去至多 s-1 个 1的子数组数量得到恰好 s 个 1 的子数组数量。算法正确性证明计数方法的正确性对于至多 s 个 1的子数组数量使用滑动窗口方法计算。每次扩展右边界统计以当前右边界结尾的满足条件的子数组数量right - left 1。然后收缩左边界直到窗口内 1 的数量不超过 s。减法原理恰好 s 个 1 的子数组 至多 s 个 1 的子数组 - 至多 s-1 个 1 的子数组。复杂度分析哈希表方法时间复杂度O(n)空间复杂度O(n)计数方法时间复杂度O(n)两次遍历空间复杂度O(1)代码实现Python 实现哈希表def numSubarraysWithSum(nums, s): prefix_sum 0 count 0 prefix_map {0: 1} for num in nums: prefix_sum num count prefix_map.get(prefix_sum - s, 0) prefix_map[prefix_sum] prefix_map.get(prefix_sum, 0) 1 return countPython 实现计数def numSubarraysWithSum_count(nums, s): def countAtMost(k): if k 0: return 0 result 0 left 0 count 0 for right in range(len(nums)): if nums[right] 1: count 1 while count k: if nums[left] 1: count - 1 left 1 result right - left 1 return result return countAtMost(s) - countAtMost(s - 1)边界情况处理S 0当 S 0 时需要统计和为 0 的子数组数量即没有 1 的子数组。计数方法中的 countAtMost(-1) 会返回 0countAtMost(0) 会正确计算没有 1 的子数组数量。全 0 数组当数组全为 0 时子数组和都是 0。如果 S 0子数组数量是 n*(n1)/2如果 S 0子数组数量是 0。两种方法都能正确处理。全 1 数组当数组全为 1 时子数组和等于子数组长度。计数方法使用滑动窗口正确统计恰好 s 个 1 的子数组数量。测试用例def test_num_subarrays_with_sum(): assert numSubarraysWithSum([1,0,1,0,1], 2) 4 assert numSubarraysWithSum([1,0,1,0,1], 0) 4 assert numSubarraysWithSum([0,0,0,0,0], 0) 15 assert numSubarraysWithSum([1,1,1,1], 2) 3 assert numSubarraysWithSum([1], 0) 0 assert numSubarraysWithSum([0], 0) 1 assert numSubarraysWithSum([], 0) 1 print(所有测试用例通过)总结和相同的二元子数组问题展示了两种方法通用的前缀和哈希表方法和针对二元数组优化的计数方法。两种方法都是 O(n) 时间复杂度但计数方法空间复杂度更优O(1) vs O(n)。对于二元数组计数方法利用了恰好 s 个 1的特性通过计算至多 s 个和至多 s-1 个的差来得到精确答案。这种转换思路在处理恰好问题时很有用。
