LeetCode-Python-3101. 交替子数组计数

给你一个

二进制数组

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