🔗 🟡 3101. Count Alternating Subarrays 1405

tags: `Weekly Contest 391` `陣列(Array)` `數學(Math)` `動態規劃(DP)`

題意

給定一個二進位制陣列 $nums$ 。

如果一個子陣列中相鄰的兩個元素值 不相同 ，我們就稱其為 交替的(alternating) 。

返回 $nums$ 中交替子陣列的數量。

思路

方法一：動態規劃(DP)

令 $dp[i]$ 表示以 $nums[i]$ 結尾的交替子陣列的數量，則有以下遞迴式：

若 $nums[i] \neq nums[i-1]$ ，則代表 $nums[i]$ 可以繼續接在 $nums[i-1]$ 後面，因此 $dp[i] = dp[i-1] + 1$ ；
否則 $dp[i] = 1$ ，因為 $nums[i]$ 本身就可以視為一個交替子陣列。

初始化 $dp$ 陣列為全 $1$ ，從 $i=1$ 開始遍歷 $nums$ 陣列，更新 $dp[i]$ 和答案 $ans$ 。由於計算時會忽略 $i=0$ 的情況，因此須將答案初始化為 $1$ 。

複雜度分析

時間複雜度： $O(n)$ ，其中 $n$ 為陣列 $nums$ 的長度。
空間複雜度： $O(n)$ ，用於存儲 $dp$ 陣列。

class Solution:
    def countAlternatingSubarrays(self, nums: List[int]) -> int:
        n = len(nums)
        ans = 1
        dp = [1] * n
        for i in range(1, n):
            if nums[i] != nums[i - 1]:
                dp[i] = dp[i - 1] + 1
            ans += dp[i]
        return ans

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

方法二：空間優化

注意到方法一中，我們只需要用到 $dp[i]$ 和 $dp[i-1]$ 的值，因此可以將空間複雜度優化至 $O(1)$ 。

令 $f$ 表示以 $nums[i]$ 結尾的交替子陣列的數量，則若 $nums[i] \neq nums[i-1]$ ，則 $f = f + 1$ ；否則 $f = 1$ ，並在計算過程中累加 $f$ 的值至答案中，同樣需將 $ans$ 初始化為 $1$ 。

複雜度分析

時間複雜度： $O(n)$ ，其中 $n$ 為陣列 $nums$ 的長度。
空間複雜度： $O(1)$ 。

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

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