给你一个
二进制数组
nums
。
如果一个
子数组
中 不存在 两个 相邻 元素的值 相同 的情况,我们称这样的子数组为 交替子数组 。
返回数组 nums
中交替子数组的数量。
示例 1:
输入: nums = [0,1,1,1]
输出: 5
解释:
以下子数组是交替子数组:[0]
、[1]
、[1]
、[1]
以及 [0,1]
。
示例 2:
输入: nums = [1,0,1,0]
输出: 10
解释:
数组的每个子数组都是交替子数组。可以统计在内的子数组共有 10 个。
提示:
1 <= nums.length <= 105
nums[i]
不是0
就是1
。
思路:
首先如果我们有一个长度为 n 的数组,并且它的交替元素都不相同,那么它的交替子数组个数应该等于 1 + 2 + 3 + .. + n, 即所有子数组的个数。因为每个子数组都是交替子数组。
那么对应这道题目,只需要把数组分成子数组,满足每个子数组内相邻元素均不相同,然后统计每个部分的解最后累加即可。
比如对于 nums = [0, 1, 1, 1], 就可以分成 [0, 1], [1], [1]。其答案就是 1 + 2 + 1 + 1 = 5
时间复杂度:O(N)
空间复杂度:O(1)
class Solution:
def countAlternatingSubarrays(self, nums: List[int]) -> int:
# divide nums into different parts where in each part, all nearby elements are diff
# in a qualified part, the number of answer is sigma(1, n)(n = part_length)
lens = []
subarray_start = res = 0
for i, num in enumerate(nums):
if i and num == nums[i - 1]:
# divide between nums[i - 1] and num, subarray = nums[subarray_start:i]
subarray_len = i - subarray_start
res += (1 + subarray_len) * (subarray_len) // 2
subarray_start = i
# print(res)
k = len(nums) - subarray_start
return res + (1 + k) * k // 2