LeetCode 🔴 3700. Number of ZigZag Arrays II

🔗 🔴 3700. Number of ZigZag Arrays II

Problem Statement

題目簡述

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

• 每個元素都在 $[l, r]$ 之間
• 相鄰元素不相等
• 任三個連續元素不能形成嚴格遞增或嚴格遞減序列

答案對 $10^9 + 7$ 取模。換句話說，陣列的相鄰大小關係必須嚴格交替，也就是 ZigZag。

Constraints

約束條件

• $3 \le n \le 10^9$
• $1 \le l < r \le 75$

---

思路：DP 線性轉移 + 矩陣快速冪

前置知識

這題是 LeetCode 🔴 3699. Number of ZigZag Arrays I 的延伸。建議先讀完 3699 的 DP 推導，尤其是「上升 / 下降交替」的狀態機，以及方法三中「嚴格前綴和後反轉」的統一轉移，接下來不再重新推導 ZigZag 的 DP 轉移。

注意兩題的差別只在資料範圍：

• 在 I 中 $n,V \le 2000$，可以逐步做 DP；
• 而在 II 中 $n$ 放大到 $10^9$，但值域縮小到 $V \le 75$。

現在問題變成：轉移輪數非常多，但狀態數非常少。而逐步 DP 的瓶頸在 $n$，不是 $V$，通常這種情況要把 $n$ 利用快速冪壓縮到 $\log n$，由於下一層的狀態是前綴和 / 後綴和，可以表示成上一層狀態的線性組合，因此可以把「從上一層 DP 到下一層 DP」寫成一個轉移矩陣，然後用矩陣快速冪計算。

方法一：拚接兩個狀態向量

回顧定義

接續 LeetCode 🔴 3699. Number of ZigZag Arrays I 的雙狀態 DP，延續之前定義的 $f_0$ 和 $f_1$。

• $f_0[i][j]$：考慮前 $i$ 個數，最後一個數字為 $j$，且最後一步是上升的方案數。
• $f_1[i][j]$：考慮前 $i$ 個數，最後一個數字為 $j$，且最後一步是下降的方案數。

初始值：

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

把狀態轉換成向量形式，並找出轉移矩陣

由於轉移其實就是把 $f_0$ 更新為 $f_1$ 的 prefix sum、把 $f_1$ 更新為 $f_0$ 的 suffix sum，兩者交替進行，所以可以把兩個方向的狀態拼成一個長度為 $2V$ 的向量，然後寫出一個 $2V \times 2V$ 的轉移矩陣。

令：

$$\mathbf f_i = \begin{bmatrix} f_0[i][0] & f_0[i][1] & \cdots & f_0[i][V-1] & f_1[i][0] & f_1[i][1] & \cdots & f_1[i][V-1] \end{bmatrix}^{\mathsf T}$$

初始向量為全 $1$：

$$\mathbf f_1 = \begin{bmatrix} 1 & 1 & \cdots & 1 & 1 & 1 & \cdots & 1 \end{bmatrix}^{\mathsf T}$$

目標是找到一個 $2V \times 2V$ 的矩陣 $M$，使得：

$$\mathbf f_i = M\mathbf f_{i-1}$$

由於狀態向量是把 $f_0$ 與 $f_1$ 拼接起來，且 $f_0$ 的狀態依賴於 $f_1$、$f_1$ 的狀態依賴於 $f_0$，因此 $M$ 可以直接寫成如下的分塊矩陣：

$$M= \begin{bmatrix} 0 & A \\ B & 0 \end{bmatrix}$$

其中：

$$A_{i,j}=[j<i],\qquad B_{i,j}=[j>i]$$

例如 $V=4$

$$M= \begin{bmatrix} 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 & 1 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 & 1 & 1 & 0 & 0 \\ 0 & 0 & 0 & 0 & 1 & 1 & 1 & 0 \\ 0 & 1 & 1 & 1 & 0 & 0 & 0 & 0 \\ 0 & 0 & 1 & 1 & 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 1 & 0 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 & 0 & 0 & 0 & 0 \end{bmatrix}$$

利用矩陣快速冪計算答案

所以：

$$\mathbf f_n = M^{n-1}\mathbf f_1$$

利用矩陣快速冪計算，答案為 $\sum_x \mathbf f_n[x]$。

複雜度分析

• 時間複雜度：$\mathcal{O}(V^3\log n)$。矩陣大小為 $2V\times 2V$，一次矩陣乘法需要 $\mathcal{O}((2V)^3)=\mathcal{O}(V^3)$；快速冪需要做 $\mathcal{O}(\log n)$ 次矩陣乘法。
• 空間複雜度：$\mathcal{O}(V^2)$。需要保存 $2V\times 2V$ 的轉移矩陣，以及長度為 $2V$ 的狀態向量。

方法二：利用對稱性優化，並統一操作方式

接續 LeetCode 🔴 3699. Number of ZigZag Arrays I 的方法二和方法三，可以利用對稱性優化，並把操作統一成：做「嚴格前綴和」之後「反轉」。

因此只需要維護一個長度為 $V$ 的狀態向量 $\mathbf f_i$，且轉移矩陣 $M$ 的大小也只需要 $V\times V$。

由於反轉後新位置 $x$ 對應原位置 $V-1-x$，所以：

$$M[x][j] = 1 \quad \Longleftrightarrow \quad j < V-x-1$$

也就是反對角線左上方全為 $1$，其餘為 $0$。

例如 $V=4$

$$M= \begin{bmatrix} 1 & 1 & 1 & 0 \\ 1 & 1 & 0 & 0 \\ 1 & 0 & 0 & 0 \\ 0 & 0 & 0 & 0 \end{bmatrix}$$

因此：

$$\mathbf f_n = M^{n-1}\mathbf f_1,\qquad \mathbf f_1=\mathbf 1$$

由於只計算了一半的答案，因此最終答案要 $\times 2$。答案為：

$$2\sum_x \mathbf f_n[x]$$

複雜度分析

• 時間複雜度：$\mathcal{O}(V^3\log n)$。矩陣大小降為 $V\times V$，一次矩陣乘法為 $\mathcal{O}(V^3)$，快速冪需要 $\mathcal{O}(\log n)$ 次乘法；相比方法一，主要是常數因子更小。
• 空間複雜度：$\mathcal{O}(V^2)$。需要保存 $V\times V$ 的轉移矩陣與長度為 $V$ 的狀態向量。

本題要點

1. ZigZag 約束可以看成上升 / 下降交替的 DP 狀態機。
2. $n$ 很大時，不能逐層 DP，要把固定轉移改寫成矩陣快速冪。
3. 利用對稱性與翻轉操作，可以把雙方向狀態壓縮成單一狀態矩陣。

---

Code

方法一：拚接兩個狀態向量

普通寫法numpy優化

MOD = int(1e9) + 7

class Matrix:
    def __init__(self, mat):
        self.mat = mat

    def __mul__(self, other):
        assert len(self.mat[0]) == len(other.mat)
        return Matrix([[sum(x * y for x, y in zip(row, col)) % MOD for col in zip(*other.mat)]
                for row in self.mat])
    
    def __pow__(self, k):
        assert len(self.mat) == len(self.mat[0])
        n = len(self.mat)
        res = Matrix([[int(i == j) for j in range(n)] for i in range(n)])
        base = self
        while k > 0:
            if k & 1:
                res *= base
            base *= base
            k >>= 1
        return res
    
    def mul_pow(self, k, x):
        assert len(self.mat[0]) == len(x.mat)
        res = x
        base = self
        while k > 0:
            if k & 1:
                res = base * res
            base *= base
            k >>= 1
        return res
    
    def __getitem__(self, i):
        return self.mat[i]

    def __str__(self):
        return "\n".join(map(lambda x: " ".join(map(str, x)), self.mat))

class Solution:
    def zigZagArrays(self, n: int, l: int, r: int) -> int:
        V = r - l + 1  # [l, r] => [0, V - 1]
        f = Matrix([[1] for _ in range(V * 2)])
        M = [[0] * (V * 2) for _ in range(V * 2)]
        for i in range(V):
            for j in range(V, V + i):
                M[i][j] = 1
            for j in range(i + 1, V):
                M[i + V][j] = 1
        M = Matrix(M)
        f = M.mul_pow(n - 1, f)
        return sum(map(sum, f.mat)) % MOD

import numpy as np

MOD = int(1e9) + 7

def matrix_power_mul(A: np.ndarray, n: int, x: np.ndarray, mod: int):
    assert len(A[0]) == len(x)
    res = x
    while n > 0:
        if n & 1:
            res = A @ res % mod
        A = A @ A % mod
        n >>= 1
    return res

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 * 2, 1), dtype=np.object_)
        M = np.zeros((V * 2, V * 2), dtype=np.object_)
        for i in range(V):
            for j in range(V, V + i):
                M[i][j] = 1
            for j in range(i + 1, V):
                M[i + V][j] = 1
        f = matrix_power_mul(M, n - 1, f, MOD)
        return int(np.sum(f) % MOD)

方法二：利用對稱性優化，並統一操作方式

普通寫法numpy優化

MOD = int(1e9) + 7

class Matrix:
    def __init__(self, mat):
        self.mat = mat

    def __mul__(self, other):
        assert len(self.mat[0]) == len(other.mat)
        return Matrix([[sum(x * y for x, y in zip(row, col)) % MOD for col in zip(*other.mat)]
                for row in self.mat])
    
    def __pow__(self, k):
        assert len(self.mat) == len(self.mat[0])
        n = len(self.mat)
        res = Matrix([[int(i == j) for j in range(n)] for i in range(n)])
        base = self
        while k > 0:
            if k & 1:
                res *= base
            base *= base
            k >>= 1
        return res
    
    def mul_pow(self, k, x):
        assert len(self.mat[0]) == len(x.mat)
        res = x
        base = self
        while k > 0:
            if k & 1:
                res = base * res
            base *= base
            k >>= 1
        return res
    
    def __getitem__(self, i):
        return self.mat[i]

    def __str__(self):
        return "\n".join(map(lambda x: " ".join(map(str, x)), self.mat))

class Solution:
    def zigZagArrays(self, n: int, l: int, r: int) -> int:
        V = r - l + 1  # [l, r] => [0, V - 1]
        f = Matrix([[1] for _ in range(V)])
        M = [[0] * V for _ in range(V)]
        for i in range(V):
            for j in range(V - i - 1):
                M[i][j] = 1
        M = Matrix(M)
        f = M.mul_pow(n - 1, f)
        return sum(map(sum, f.mat)) * 2 % MOD

import numpy as np

MOD = int(1e9) + 7

def matrix_power_mul(A: np.ndarray, n: int, x: np.ndarray, mod: int):
    assert len(A[0]) == len(x)
    res = x
    while n > 0:
        if n & 1:
            res = A @ res % mod
        A = A @ A % mod
        n >>= 1
    return res

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, 1), dtype=np.object_)
        M = np.tril(np.ones((V, V), dtype=np.object_)[::-1], k=-1)[::-1]
        f = matrix_power_mul(M, n - 1, f, MOD)
        return int(np.sum(f) * 2 % MOD)