3101. 交替子数组计数

English Version

题目描述

给你一个二进制数组 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 。

解法

方法一：枚举

我们可以枚举以每个位置结尾的子数组，计算满足条件的子数组的数量，累加所有位置的满足条件的子数组的数量即可。

具体地，我们定义变量 $s$ 表示以元素 $nums[i]$ 结尾的满足条件的子数组的数量，初始时我们将 $s$ 置为 $1$，表示以第一个元素结尾的满足条件的子数组的数量为 $1$。

接下来，我们从第二个元素开始遍历数组，对于每个位置 $i$，我们根据 $nums[i]$ 和 $nums[i-1]$ 的关系更新 $s$ 的值：

• 如果 $nums[i] \neq nums[i-1]$，则 $s$ 的值增加 $1$，即 $s = s + 1$；
• 如果 $nums[i] = nums[i-1]$，则 $s$ 的值重置为 $1$，即 $s = 1$。

然后，我们将 $s$ 的值累加到答案中，继续遍历数组的下一个位置，直到遍历完整个数组。

时间复杂度 $O(n)$，其中 $n$ 是数组的长度。空间复杂度 $O(1)$。

class Solution:
    def countAlternatingSubarrays(self, nums: List[int]) -> int:
        ans = s = 1
        for a, b in pairwise(nums):
            s = s + 1 if a != b else 1
            ans += s
        return ans

class Solution {
    public long countAlternatingSubarrays(int[] nums) {
        long ans = 1, s = 1;
        for (int i = 1; i < nums.length; ++i) {
            s = nums[i] != nums[i - 1] ? s + 1 : 1;
            ans += s;
        }
        return ans;
    }
}

class Solution {
public:
    long long countAlternatingSubarrays(vector<int>& nums) {
        long long ans = 1, s = 1;
        for (int i = 1; i < nums.size(); ++i) {
            s = nums[i] != nums[i - 1] ? s + 1 : 1;
            ans += s;
        }
        return ans;
    }
};

func countAlternatingSubarrays(nums []int) int64 {
	ans, s := int64(1), int64(1)
	for i, x := range nums[1:] {
		if x != nums[i] {
			s++
		} else {
			s = 1
		}
		ans += s
	}
	return ans
}

function countAlternatingSubarrays(nums: number[]): number {
    let [ans, s] = [1, 1];
    for (let i = 1; i < nums.length; ++i) {
        s = nums[i] !== nums[i - 1] ? s + 1 : 1;
        ans += s;
    }
    return ans;
}