和相同的二元子数组

标签: 数组 哈希表 前缀和 滑动窗口

难度: Medium

给你一个二元数组 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 * 104
  • nums[i] 不是 0 就是 1
  • 0 <= goal <= nums.length

Submission

运行时间: 206 ms

内存: 16.8 MB

class Solution:
    def numSubarraysWithSum(self, nums: List[int], goal: int) -> int:
        res = 0
        cnt = 0

        if goal == 0:
            for n in nums:
                if n == 1:
                    cnt = 0
                else:
                    cnt += 1
                res += cnt
        else:
            i = 0
            for n in nums:
                if n == 1:
                    goal -= 1
                    cnt = 0
                    while goal == 0:
                        if nums[i] == 1:
                            goal += 1
                        i += 1
                        cnt += 1
                res += cnt

        return res

Explain

该题解利用滑动窗口的思想处理子数组问题。对于goal为0的特殊情况,单独处理,统计连续的0的子数组个数。对于goal不为0的情况,遍历数组,每遇到1,则减少goal的值,如果goal减为0,则说明从当前窗口的开始到当前位置的子数组是一个有效的子数组。通过调整窗口的起始位置i,直到窗口内的和再次大于goal,过程中统计每次有效的子数组个数。这种方法只需要遍历数组一次,同时调整窗口起始位置,直到整个数组都被检查过。

时间复杂度: O(n)

空间复杂度: O(1)

class Solution:
    def numSubarraysWithSum(self, nums: List[int], goal: int) -> int:
        res = 0  # 结果计数器
        cnt = 0  # 当前窗口内满足条件的子数组数量

        if goal == 0:
            for n in nums:
                if n == 1:
                    cnt = 0  # 遇到1时重置计数器
                else:
                    cnt += 1  # 增加连续0的计数
                res += cnt  # 累加结果
        else:
            i = 0  # 窗口起始位置
            for n in nums:
                if n == 1:
                    goal -= 1  # 遇到1时减少goal
                    cnt = 0  # 重置当前窗口内满足条件的子数组计数
                    while goal == 0:
                        if nums[i] == 1:
                            goal += 1  # 调整窗口直至窗口内的和大于goal
                        i += 1  # 移动窗口起始位置
                        cnt += 1  # 增加计数
                res += cnt  # 累加满足条件的子数组个数

        return res

Explore

在处理goal为0的情况时,算法通过连续计数从当前位置开始到目前为止的连续0的个数来确保不漏计。每遇到一个0,计数器cnt加1,表示新的以当前0结尾的子数组。同时,每次遇到0时,将当前cnt的值累加到结果res中,这样每个包含0的子数组都会被计算在内,包括长度为1的子数组,以及所有可能的更长的以当前0结尾的子数组。当遇到1时,计数器重置为0,因为以1结尾的子数组不满足和为0的要求。

在goal不为0的情况下,每当遇到1时,goal的值减小,并且子数组计数器cnt重置为0,因为需要重新计算新的可能的子数组。然后,如果当前窗口内的和等于goal(goal减到0),则开始从窗口开始位置i进行调整,一直调整到窗口内的和再次大于goal。在这个过程中,每次调整都增加计数器cnt,累加到结果res中。这种方法确保了每个符合条件的子数组只被计算一次,因为每次遇到1时,我们都会从新的子数组开始计数,避免重复计数。

在调整窗口起始位置i时,选择在`goal == 0`的条件下进行调整是因为这表示窗口内的子数组和已经达到了目标goal,是一个有效的子数组。这种调整的特别考虑是为了找到所有可能的以当前窗口结尾的子数组,满足和为goal。通过逐渐移动窗口的起始位置i,减小窗口的大小,直到窗口内的和大于goal,我们可以找出所有包括当前窗口结束位置的符合条件的子数组。这种方法可以确保我们不会漏掉任何一个符合条件的子数组,同时也避免了重复计数。