题目
给定一个二进制数组 n u m s nums nums,找到含有相同数量的0和1的最长连续子数组,并返回该子数组的长度。
示例:
输入:nums = [0,1,0]
输出:2
说明:[0,1](或[1,0])是具有相同数量0和1的最长连续子数组。
方法:前缀和+哈希表
由于“0和1数量相同”等价于“1的数量减去0的数量等于0”,可以将数组中的0视作-1,则原问题转换成“求最长的连续子数组,其元素和为0”。
设数组
n
u
m
s
nums
nums的长度为
n
n
n,将数组
n
u
m
s
nums
nums进行转换得到长度相等的新数组
n
e
w
N
u
m
s
newNums
newNums:对于
0
≤
i
<
n
0\leq i <n
0≤i<n,当
n
u
m
s
[
i
]
=
1
nums[i] = 1
nums[i]=1时
n
e
w
N
u
m
s
[
i
]
=
1
newNums[i] =1
newNums[i]=1,当
n
u
m
s
[
i
]
=
0
nums[i]=0
nums[i]=0时
n
e
w
N
u
m
s
[
i
]
=
−
1
newNums[i]=-1
newNums[i]=−1。
为了快速计算
n
e
w
N
u
m
s
newNums
newNums的子数组的元素和,需要首先计算
n
e
w
N
u
m
s
newNums
newNums的前缀和。用
p
r
e
f
i
x
S
u
m
s
[
i
]
prefixSums[i]
prefixSums[i]表示
n
e
w
N
u
m
s
newNums
newNums从下标0到下标
i
i
i的前缀和,则
n
e
w
N
u
m
s
newNums
newNums从下标
j
+
1
j+1
j+1到下标
k
k
k(其中
j
<
k
j<k
j<k)的子数组元素和为
p
r
e
f
i
x
S
u
m
s
[
k
]
−
p
r
e
f
i
x
S
u
m
s
[
j
]
prefixSums[k]-prefixSums[j]
prefixSums[k]−prefixSums[j],该子数组的长度为
k
−
j
k-j
k−j。
当
p
r
e
f
i
x
S
u
m
s
[
k
]
−
p
r
e
f
i
x
S
u
m
s
[
j
]
=
0
prefixSums[k]-prefixSums[j]=0
prefixSums[k]−prefixSums[j]=0时,即得到
n
e
w
N
u
m
s
newNums
newNums的一个长度为
k
−
j
k-j
k−j的子数组元素和为0,对应
n
u
m
s
nums
nums的一个长度为
k
−
j
k-j
k−j的子数组中有相同数量的0和1。
实现方面,不需要创建数组
n
e
w
N
u
m
s
newNums
newNums和
p
r
e
f
i
x
S
u
m
s
prefixSums
prefixSums,只需要维护一个变量
c
o
u
n
t
e
r
counter
counter存储
n
e
w
N
u
m
s
newNums
newNums的前缀和即可。具体做法是,遍历数组
n
u
m
s
nums
nums,当遇到元素1时将
c
o
u
n
t
e
r
counter
counter的值加1,当遇到元素0时将
c
o
u
n
t
e
r
counter
counter的值减1,遍历过程中使用哈希表存储每个前缀和第一次出现的下标。
规定空的前缀的结束下标为-1,由于空的前缀的元素和为0,因此在遍历之前,首先在哈希表中存入键值对(0,-1)。遍历过程中,对于每个下标
i
i
i,进行如下操作:
- 如果 c o u n t e r counter counter的值在哈希表中已经存在,则取出 c o u n t e r counter counter在哈希表中对应的下标 p r e v I n d e x , n u m s prevIndex,nums prevIndex,nums从下标 p r e v I n d e x + 1 prevIndex+1 prevIndex+1到下标 i i i的子数组中有相同数量的0和1,该子数组的长度为 i − p r e v I n d e x i-prevIndex i−prevIndex,使用该子数组的长度更新最长连续子数组的长度;
- 如果 c o u n t e r counter counter的值在哈希表中不存在,则将当前余数和当前下标 i i i的键值对存入哈希表中。
由于哈希表存储的是 c o u n t e r counter counter的每个取值第一次出现的下标,因此当遇到重复的前缀和时,根据当前下标和哈希表中存储的下标计算得到的子数组长度是以当前下标结尾的子数组中满足有相同数量的0和1的最长子数组的长度。遍历结束时,即可得到 n u m s nums nums中的有相同数量的0和1的最长子数组的长度。
代码
class Solution {
public:
int findMaxLength(vector<int>& nums) {
int maxLength = 0;
unordered_map<int, int> m;
int counter = 0;
m[counter] = -1;
for(int i = 0; i < nums.size(); ++i) {
int num = nums[i];
if(num == 1) {
++counter;
} else {
--counter;
}
if(m.count(counter)) {
maxLength = max(maxLength, i - m[counter]);
} else {
m[counter] = i;
}
}
return maxLength;
}
};