3700. Z字形数组II

力扣

力扣杯 469

3700. Z字形数组II

题目描述

给你三个整数 n、l 和 r。

长度为 n 的 Z字形数组 定义如下：

每个元素都在 [l, r] 范围内。
没有两个相邻的元素是相等的。
没有三个连续的元素形成严格递增或严格递减的序列。

返回有效Z字形数组的总数。由于答案可能很大，请将其对 $10^9 + 7$ 取模后返回。

此问题是 3699. Z字形数组I 的延续。唯一的区别在于 n 的约束条件，现在 n 最大可达 $10^9$ ，而范围 r - l 最多为 75。

解题思路

此问题直接建立在 Z字形数组I 的动态规划解法之上。核心的 DP 状态和转移保持不变，但 n 的值大得多，使得线性时间的 DP 方法过慢。我们必须找到一种方法来优化第 n 个 DP 状态的计算。

从 DP 到矩阵形式

让我们回顾一下 DP 的转移。在将范围 [l, r] 归一化为 [1, m]（其中 m = r - l + 1）之后，我们有：

dp[i][j][0] (以 ... < j 结尾): $dp[i][j][0] = \sum_{p=j+1}^{m} dp[i-1][p][1]$
dp[i][j][1] (以 ... > j 结尾): $dp[i][j][1] = \sum_{p=1}^{j-1} dp[i-1][p][0]$

这个线性递推系统强烈暗示我们可以使用矩阵。让我们定义一个大小为 2m 的状态向量 V_i，它表示给定长度 i 的 DP 值： V_i = [dp[i][1][0], ..., dp[i][m][0], dp[i][1][1], ..., dp[i][m][1]]^T

我们的目标是找到一个 2m x 2m 的转移矩阵 M，使得 V_i = M * V_{i-1}。观察依赖关系，dp[i][...][0] 只依赖于 dp[i-1][...][1]，而 dp[i][...][1] 只依赖于 dp[i-1][...][0]。这种结构表明 M 是一个分块矩阵：

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

这里，0 是一个 m x m 的零矩阵，A 和 B 是处理转移的 m x m 矩阵。

矩阵 A: 这个矩阵将“向下”状态 (dp[i-1][...][1]) 转换为“向上”状态 (dp[i][...][0])。方程是 dp[i][j][0] = sum_{p=j+1 to m} dp[i-1][p][1]。这意味着 A 的第 j 行在所有 p > j 的列上都应该是 1。 $A_{jk} = \begin{cases} 1 & \text{if } k > j \\ 0 & \text{otherwise} \end{cases}$
矩阵 B: 这个矩阵将“向上”状态 (dp[i-1][...][0]) 转换为“向下”状态 (dp[i][...][1])。方程是 dp[i][j][1] = sum_{p=1 to j-1} dp[i-1][p][0]。这意味着 B 的第 j 行在所有 p < j 的列上都应该是 1。 $B_{jk} = \begin{cases} 1 & \text{if } k < j \\ 0 & \text{otherwise} \end{cases}$

例如，如果 m = 4，矩阵 A 和 B 将是： $A = \begin{pmatrix} 0 & 1 & 1 & 1 \\ 0 & 0 & 1 & 1 \\ 0 & 0 & 0 & 1 \\ 0 & 0 & 0 & 0 \end{pmatrix}, \quad B = \begin{pmatrix} 0 & 0 & 0 & 0 \\ 1 & 0 & 0 & 0 \\ 1 & 1 & 0 & 0 \\ 1 & 1 & 1 & 0 \end{pmatrix}$

矩阵快速幂

有了转移 V_i = M * V_{i-1}，我们可以通过从一个基本情况开始并重复应用矩阵 M 来找到 V_n。 V_n = M * V_{n-1} = M^2 * V_{n-2} = ... = M^{n-2} * V_2

由于 n 可能非常大，我们不能执行 n-2 次矩阵乘法。相反，我们使用矩阵快速幂（也称为矩阵的二进制幂）在 $O((2m)^3 \log n)$ 时间内计算 $M^{n-2}$ 。

基本情况是状态向量 V_2，它代表长度为 2 的数组。

dp[2][j][0] (A[1] < A[2]=j): A[1] 可以是从 l 到 l+j-2 的任何值，所以有 j-1 种可能性。
dp[2][j][1] (A[1] > A[2]=j): A[1] 可以是从 l+j 到 r 的任何值，所以有 m-j 种可能性。

所以，初始状态向量 V_2 是： V_2 = [m-1, m-2, ..., 0, 0, 1, ..., m-1]^T （注意：提供的代码在 DP 基本情况中有一个轻微的偏差，将 [1,m] 映射到 dp[2][j][0] = m-j 和 dp[2][j][1] = j-1。逻辑保持不变。）

在计算出 $M' = M^{n-2}$ 之后，我们计算最终状态 V_n = M' * V_2。Z字形数组的总数是 V_n 中所有元素的和，对 $10^9 + 7$ 取模。

查看代码

static const auto fast_io = []() {
    std::ios_base::sync_with_stdio(false);
    std::cin.tie(NULL);
    return 0;
}();
#define N 2005
class Solution {
public:
    int mod = 1e9 + 7;

    vector<vector<long long>> multi(vector<vector<long long>>& a,
                                    vector<vector<long long>>& b) {
        int n = a.size();
        vector<vector<long long>> c(n, vector<long long>(n, 0));
        for (int i = 0; i < n; i++) {
            for (int j = 0; j < n; j++) {
                for (int k = 0; k < n; k++) {
                    c[i][j] = (c[i][j] + a[i][k] * b[k][j]) % mod;
                }
            }
        }
        return c;
    }

    vector<vector<long long>> power(vector<vector<long long>> a, int n) {
        int len = a.size();
        vector<vector<long long>> res(len, vector<long long>(len));
        for (int i = 0; i < len; i++)
            res[i][i] = 1;
        while (n) {
            if (n & 1)
                res = multi(res, a);
            a = multi(a, a);
            n /= 2;
        }
        return res;
    }

    int zigZagArrays(int n, int l, int r) {
        if (n == 1) {
            return r - l + 1;
        }
        if (n == 2) {
            long long m = r - l + 1;
            return (m * (m - 1)) % mod;
        }
        long long m = r - l + 1;
        vector<vector<long long>> tran(2 * m, vector<long long>(2 * m, 0));
        for (int i = 0; i < m; i++) {
            for (int j = i + 1; j < m; j++) {
                tran[i][j + m] = 1;
            }
        }
        for (int i = 1; i < m; i++) {
            for (int j = 0; j < i; j++) {
                tran[i + m][j] = 1;
            }
        }
        tran = power(tran, n - 2);
        vector<long long> start(2 * m, 0);
        for (int i = 0; i < m; i++) {
            start[i] = m - 1 - i;
            start[i + m] = i;
        }
        vector<long long> end(2 * m, 0);
        for (int i = 0; i < 2 * m; i++) {
            for (int j = 0; j < 2 * m; j++) {
                end[i] = (end[i] + tran[i][j] * start[j]) % mod;
            }
        }
        long long ans = 0;
        for (int i = 0; i < 2 * m; i++) {
            ans = (ans + end[i]) % mod;
        }
        return ans;
    }
};

时间和空间复杂度

时间复杂度: $O((r-l)^3 \log n)$ 。矩阵的大小是 2m x 2m，其中 m = r-l+1。矩阵乘法需要 $O(m^3)$ ，而快速幂需要 $O(\log n)$ 次乘法。
空间复杂度: $O((r-l)^2)$ 用于存储转移矩阵。

3699. Z字形数组I