剑指 Offer II 011. 0 和 1 个数相同的子数组

• 标签：数组、哈希表、前缀和
• 难度：中等

题目链接

• 剑指 Offer II 011. 0 和 1 个数相同的子数组 - 力扣

题目大意

给定一个二进制数组 nums。

要求：找到含有相同数量 0 和 1 的最长连续子数组，并返回该子数组的长度。

解题思路

「0 和 1 数量相同」等价于「1 的数量减去 0 的数量等于 0」。

我们可以使用一个变量 pre_diff 来记录下前 i 个数中，1 的数量比 0 的数量多多少个。我们把这个 pre_diff叫做「1 和 0 数量差」，也可以理解为变种的前缀和。

然后我们再用一个哈希表 pre_dic 来记录「1 和 0 数量差」第一次出现的下标。

那么，如果我们在遍历的时候，发现 pre_diff 相同的数量差已经在之前出现过了，则说明：这两段之间相减的 1 和 0 数量差为 0。

什么意思呢？

比如说：j < i，前 j 个数中第一次出现 pre_diff == 2 ，然后前 i 个数中个第二次又出现了 pre_diff == 2。那么这两段形成的子数组 nums[j + 1: i] 中 1 比 0 多 0 个，则 0 和 1 数量相同的子数组长度为 i - j。

而第二次之所以又出现 pre_diff == 2 ，是因为前半段子数组 nums[0: j] 贡献了相同的差值。

接下来还有一个小问题，如何计算「1 和 0 数量差」？

我们可以把数组中的 1 记为贡献 +1，0 记为贡献 -1。然后使用一个变量 count，只要出现 1 就让 count 加上 1，意思是又多出了 1 个 1。只要出现 0，将让 count 减去 1，意思是 0 和之前累积的 1 个 1 相互抵消掉了。这样遍历完数组，也就计算出了对应的「1 和 0 数量差」。

整个思路的具体做法如下：

• 创建一个哈希表，键值对关系为「1 和 0 的数量差：最早出现的下标 i」。
• 使用变量 pre_diff 来计算「1 和 0 数量差」，使用变量 count 来记录 0 和 1 数量相同的连续子数组的最长长度，然后遍历整个数组。
• 如果 nums[i] == 1，则让 pre_diff += 1；如果 nums[i] == 0，则让 pre_diff -= 1。
• 如果在哈希表中发现了相同的 pre_diff，则计算相应的子数组长度，与 count 进行比较并更新 count 值。
• 如果在哈希表中没有发现相同的 pre_diff，则在哈希表中记录下第一次出现 pre_diff 的下标 i。
• 最后遍历完输出 count。

注意：初始化哈希表为：pre_dic = {0: -1}，意思为空数组时，默认「1 和 0 数量差」为 0，且第一次出现的下标为 -1。

之所以这样做，是因为在遍历过程中可能会直接出现 pre_diff == 0 的情况，这种情况下说明 nums[0: i] 中 0 和 1 数量相同，如果像上边这样初始化后，就可以直接计算出此时子数组长度为 i - (-1) = i + 1。

代码

class Solution:
    def findMaxLength(self, nums: List[int]) -> int:
        pre_dic = {0: -1}
        count = 0
        pre_sum = 0
        for i in range(len(nums)):
            if nums[i]:
                pre_sum += 1
            else:
                pre_sum -= 1
            if pre_sum in pre_dic:
                count = max(count, i - pre_dic[pre_sum])
            else:
                pre_dic[pre_sum] = i
        return count