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 <= 10^5
  • nums[i] 不是 0 就是 1 。

思路

返回二进制数组中交替子数组的数量。

记录相邻元素相同的下标,计算子数组的数量。拥有n个元素的数组,其子数组数量为 n * (n + 1) / 2

官网题解是另一种思路:

  • 如果相邻的元素 nums[i-1] ≠ nums[i],那么可以将 nums[i] 加到所有以 i-1 为右端点的子数组末尾,再加上nums[i]自身,即 以 i 为右端点的交替子数组数量 = 以 i-1 为右端点的交替子数组数量 + 1
  • 如果相邻元素相等,则记录以i为右端点的交替子数组数量为1。

其实本质都是一样的,一个是使用公式求解,一个是通过循环累加。

代码

/**
 * @date 2024-07-06 20:52
 */
public class CountAlternatingSubarrays3101 {
    /**
     * 官网题解
     */
    public long countAlternatingSubarrays_v1(int[] nums) {
        long res = 0, cur = 0;
        int pre = -1;
        for (int a : nums) {
            cur = (pre != a) ? cur + 1 : 1;
            pre = a;
            res += cur;
        }
        return res;
    }

    public long countAlternatingSubarrays(int[] nums) {
        int n = nums.length;
        long res = 0;
        int s = 0;
        for (int i = 1; i < n; i++) {
            if (nums[i] == nums[i - 1]) {
                // 这里个数的计算是从 s ~ i-1的元素个数 i - 1 - s + 1
                int k = i - s;
                res += k * (k + 1) / 2;
                s = i;
            }
        }
        // 注意处理结尾的情况
        if (s != n - 1) {
            // 这里计算的是 s ~ n-1 的元素个数 n - 1 - s + 1
            int k = n - s;
            // 这里要防止溢出,使用1L
            res += k * (k + 1L) / 2;
        } else {
            res++;
        }
        return res;
    }
}

性能