跳至主要內容

0525. 连续数组

ITCharge大约 3 分钟

0525. 连续数组open in new window

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

题目链接

题目大意

给定一个二进制数组 nums

要求:找到含有相同数量 01 的最长连续子数组,并返回该子数组的长度。

解题思路

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

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

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

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

什么意思呢?

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

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

接下来还有一个小问题,如何计算「10 数量差」?

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

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

  • 创建一个哈希表,键值对关系为「10 的数量差:最早出现的下标 i」。
  • 使用变量 pre_diff 来计算「10 数量差」,使用变量 count 来记录 01 数量相同的连续子数组的最长长度,然后遍历整个数组。
  • 如果 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},意思为空数组时,默认「10 数量差」为 0,且第一次出现的下标为 -1

之所以这样做,是因为在遍历过程中可能会直接出现 pre_diff == 0 的情况,这种情况下说明 nums[0: i]01 数量相同,如果像上边这样初始化后,就可以直接计算出此时子数组长度为 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