🔗 🔴 3699. Number of ZigZag Arrays I

Problem Statement

題目簡述

給定 $n$ 、 $l$ 、 $r$ ，計算有多少個長度為 $n$ 的陣列滿足：

每個元素 $\in [l, r]$
相鄰元素不相等
任三個連續元素不成嚴格遞增或嚴格遞減

答案對 $10^9 + 7$ 取模。本質上要求陣列嚴格交替升降（ZigZag）。

Constraints

約束條件

$3 \le n \le 2000$
$1 \le l < r \le 2000$

思路：DP 狀態機 + 前綴和優化

由於值域 $[l, r]$ 的具體範圍不影響轉移邏輯，只需知道有多少個可選值。為方便討論，以下將值域平移至 $[0, V-1]$ ，其中 $V = r - l + 1$ ；之後所有狀態與轉移都在這個平移後的範圍上進行，答案不受影響。

題型定位

這題屬於「根據上一步的大小關係決定下一步限制」的 DP 狀態機模型。核心考點是如何用前綴和把 $\mathcal{O}(V^2)$ 的轉移壓到 $\mathcal{O}(V)$ ，並利用對稱性進一步簡化實作。

可遷移模型

當 DP 轉移是「對所有滿足某個不等式的先前位置求和」時，優先考慮前綴和 / 後綴和優化。對稱的轉移規則往往可以壓縮成單一狀態陣列，透過翻轉操作統一處理。

從暴力出發

最暴力的想法是：考慮到第 $i$ 個數字時，記住前一個數字與目前數字。設前一個數字為 $prev$ 、目前數字為 $curr$ ，如果要再接下一個數字 $next$ ，需要滿足：

$next \ne curr$
$(prev, curr, next)$ 不能嚴格遞增
$(prev, curr, next)$ 不能嚴格遞減

也就是說，若 $prev < curr$ ，下一步只能選 $next < curr$ ；若 $prev > curr$ ，下一步只能選 $next > curr$ 。

這樣可以定義一個很直觀的狀態：長度到 $i$ 、倒數第二個值為 $prev$ 、最後一個值為 $curr$ 的方案數。轉移時再枚舉下一個值，複雜度會是 $\mathcal{O}(nV^3)$ ，因為每一層有 $\mathcal{O}(V^2)$ 個狀態，每個狀態還要枚舉 $\mathcal{O}(V)$ 個下一個值。

觀察這個暴力狀態，其實不需要記住完整的 $prev$ 。真正影響下一步的是最後一步的大小關係：上一段是上升，下一段就必須下降；上一段是下降，下一段就必須上升。

因此引入方向維度，將狀態拆為兩種。先不做滾動優化，定義：

f_0[i][j] = \text{考慮前 } i \text{ 個數，最後一個數字為 } j\text{，且最後一步是上升的方案數}

f_1[i][j] = \text{考慮前 } i \text{ 個數，最後一個數字為 } j\text{，且最後一步是下降的方案數}

也就是說， $f_0$ 對應 $prev < curr$ ，下一步必須下降； $f_1$ 對應 $prev > curr$ ，下一步必須上升。

狀態機視角

上升與下降交替出現，形成一個兩狀態的狀態機。每一步根據當前狀態決定下一個值的合法範圍，然後翻轉狀態。

初始時長度為 $1$ ，沒有「上一步」，為了讓後續轉移統一處理，可以把兩種狀態都初始化為 $1$ ：

f_0[1][j] = f_1[1][j] = 1

這裡的兩個 $1$ 不是表示真的已經有上升或下降，而是把第一個數字作為兩種方向的共同起點。

轉移時，若要得到 $f_0[i][curr]$ ，表示最後一步要從較小的 $prev$ 走到 $curr$ ，因此前一段必須是下降結尾，累加所有 $prev < curr$ 的 $f_1[i-1][prev]$ ：

f_0[i][curr] = \sum_{prev < curr} f_1[i-1][prev]

同理：

f_1[i][curr] = \sum_{prev > curr} f_0[i-1][prev]

直接計算每次轉移是 $\mathcal{O}(V^2)$ ，總複雜度 $\mathcal{O}(nV^2)$ ，在 $n, V \le 2000$ 時無法通過。

方法一：前綴和優化 DP

注意兩個求和都是連續區間：

$\sum_{prev < curr} f_1[i-1][prev]$ 是前綴和
$\sum_{prev > curr} f_0[i-1][prev]$ 是後綴和
因此可以預處理上一層的前綴和陣列，把每個位置的轉移降為 $\mathcal{O}(1)$ 。

具體來說，令上一層兩個狀態的前綴和為：

\begin{aligned} s_0[k] &= \sum_{0 \le x < k} f_0[i-1][x] \\ s_1[k] &= \sum_{0 \le x < k} f_1[i-1][x] \end{aligned}

則轉移可以寫成：

\begin{aligned} f_0[i][curr] &= s_1[curr] \\ f_1[i][curr] &= s_0[V] - s_0[curr+1] \end{aligned}

由於第 $i$ 層只依賴第 $i-1$ 層，實作時可以滾動掉第一維，只保留上一層狀態。

複雜度分析

時間複雜度： $\mathcal{O}(nV)$
空間複雜度： $\mathcal{O}(V)$ ，只需保留上一輪的兩個狀態陣列

方法二：對稱性壓縮

關鍵觀察

注意到 $f_0$ 和 $f_1$ 的初始狀態完全相同（都是全 $1$ ），且兩者的轉移互為鏡像（一个取前綴和、一个取後綴和）。因此每一輪更新後，上升陣列和下降陣列會互為反向。

這意味著我們可以利用對稱性，實際上只需要維護一個陣列。根據當前步數的奇偶性決定執行前綴和還是後綴和：

奇數步（第 $1, 3, 5, \dots$ 次轉移）：計算前綴和，對應走到「上升結尾」
偶數步（第 $2, 4, 6, \dots$ 次轉移）：計算後綴和，對應走到「下降結尾」

最後答案乘以 $2$ 即可。

複雜度分析

時間複雜度： $\mathcal{O}(nV)$
空間複雜度： $\mathcal{O}(V)$ ，只需一個狀態陣列

方法三：翻轉統一

方法二在奇偶步之間切換前綴和與後綴和的計算方式，程式碼中有 if i & 1 的條件分支，這很麻煩，也不利於我們在 LeetCode 🔴 3700. Number of ZigZag Arrays II 中進行矩陣快速冪優化。

因為對陣列 $a$ 做後綴和（從某位置到末尾的和），等價於先將 $a$ 反轉、算前綴和、再反轉回來，因此若在 每一步 算完前綴和後都順手做一次反轉，那麼下一步的「先反轉」就剛好被上一步的「後反轉」抵消。

最終每一步的運算可以統一為：

計算前綴和
反轉陣列

等價性說明

設 $P$ 為前綴和操作、 $R$ 為反轉操作。方法二中奇數步做 $P$ 、偶數步做 $R \circ P \circ R$ （後綴和）。
方法三每步做 $R \circ P$ 。連續兩步為 $(R \circ P) \circ (R \circ P) = R \circ (P \circ R) \circ P = R \circ (R \circ \text{suffix}) \circ P$ 。經過整理，方法三的兩步等價於方法二的一奇一偶兩步加上最外層的 $R$ ，而最終求和不受反轉影響。

複雜度分析

時間複雜度： $\mathcal{O}(nV)$
空間複雜度： $\mathcal{O}(V)$

本題要點

用狀態機 DP 刻畫 ZigZag 的交替約束
利用前綴和 / 後綴和優化連續區間求和
利用對稱性合併兩個狀態陣列，再透過翻轉統一雙向操作

Code

方法一：雙狀態 DP + 前綴和

MOD = int(1e9) + 7

class Solution:
    def zigZagArrays(self, n: int, l: int, r: int) -> int:
        V = r - l + 1  # [l, r] => [0, V - 1]
        f0 = [1] * V
        f1 = [1] * V
        for _ in range(1, n):
            # 這裡加上 func=lambda x, y: (x + y) % MOD 反而會 TLE
            s0 = list(accumulate(f0, initial=0))
            s1 = list(accumulate(f1, initial=0))
            for curr in range(V):
                # 可以用 prefix sum 來優化以下這個被註解掉的迴圈
                # 且因為 f 被保存在 s 中了，所以直接覆蓋 f 就好
                # for prev in range(V):
                #     nf0[curr] += f1[prev] if curr < prev else 0
                #     nf1[curr] += f0[prev] if curr > prev else 0
                f0[curr] = s1[curr] % MOD
                f1[curr] = (s0[V] - s0[curr + 1]) % MOD
        return (sum(f0) + sum(f1)) % MOD

方法二：對稱性壓縮

MOD = int(1e9) + 7

class Solution:
    def zigZagArrays(self, n: int, l: int, r: int) -> int:
        V = r - l + 1  # [l, r] => [0, V - 1]
        f = [1] * V
        for i in range(1, n):
            # 由於 ZigZag 是有對稱性的，因此只需要算 f0 或 f1 的其中一個即可，最後再乘以 2
            # 而上述轉移其實就是把 f0 更新為 f1 的 prefix sum，f1 更新為 f0 的 suffix sum，兩者交替
            # 因此可以只保留 f0，並根據 i 的奇偶性來決定是算 prefix sum 還是 suffix sum
            if i & 1:
                f = list(accumulate(f, initial=0))[:-1]
            else:
                f = list(accumulate(f[::-1], initial=0))[:-1][::-1]
            f = [v % MOD for v in f]
        return sum(f) * 2 % MOD

import numpy as np

MOD = int(1e9) + 7

class Solution:
    def zigZagArrays(self, n: int, l: int, r: int) -> int:
        V = r - l + 1  # [l, r] => [0, V - 1]
        f = np.ones(V, dtype=np.int64)
        for i in range(1, n):
            if i & 1:
                f[1:] = np.cumsum(f[:-1])
                f[0] = 0
            else:
                f[:-1] = np.cumsum(f[:0:-1])[::-1]
                f[-1] = 0
            f = f % MOD
        return int(np.sum(f) * 2 % MOD)

方法三：翻轉統一

MOD = int(1e9) + 7

class Solution:
    def zigZagArrays(self, n: int, l: int, r: int) -> int:
        V = r - l + 1  # [l, r] => [0, V - 1]
        f = [1] * V
        for _ in range(1, n):
            # 把 i & 1 的情況中計算後綴前需要的翻轉，移動到 i & 1 == 0 中計算完成後
            # 這樣可以把兩種情況寫成同一個式子
            f = list(accumulate(f, initial=0))[:-1][::-1]
            f = [v % MOD for v in f]
        return sum(f) * 2 % MOD

import numpy as np

MOD = int(1e9) + 7

class Solution:
    def zigZagArrays(self, n: int, l: int, r: int) -> int:
        V = r - l + 1  # [l, r] => [0, V - 1]
        f = np.ones(V, dtype=np.int64)
        for _ in range(1, n):
            f[1:] = np.cumsum(f[:-1])
            f[0] = 0
            f = f[::-1] % MOD
        return int(np.sum(f)) * 2 % MOD