【LeetCode第362场周赛】8020.字符串转换 | 推导+矩阵快速幂+KMP | 困难

题目内容

原题链接

给定两个长度均为 $n$ 的字符串 $s$ 和 $t$ 。

每次选择一个长度为 $l(0<l<n)$ 的 $s$ 的后缀字符串 $suf$，将 $suif$ 从 $s$ 中删除，并添加到 $s$ 的最前面。

即 $s\leftarrow s[0:n-l]+s[n-l:n]$

现在要求这么操作恰好 $m$ 次，问有多少种不同的操作方案。

答案对 $10^9+7$ 取模。

数据范围

• $2\le n\le 5\cdot 10^5$
• $1\le m\le 10^{15}$
• $s$ 和 $t$ 均只包含小写字母

题解

剖析操作

首先可以手动模拟一些例子，可以发现无论操作多少次，得到的串必然是 $s$ 的一个长度为 $n$ 的循环串。

令 $ss=s+s$

即操作若干次后的 $ops=ss[i:i+n](0\le i<n)$ 。

动态规划

状态定义

定义 $f[i][j]$ 表示操作了 $i$ 次后，串变成了 $ss[j:j+n]$ 的方案数。

状态转移

$f[i][j]=\sum _{k=0\&k\ne j}^{n-1}f[i-1][k]$

初始化 $f[0][0]=1,f[0][j]=0(1\le j<n)$

假设有 $cnt$ 个下标 $idx$ ，满足 $ss[idx:idx+n]=t$ ，则最终的答案为：$\sum _{ss[idx:idx+n]=t}f[k][idx]$。

但这样转移显然是不满足空间时间的要求的，考虑如何优化。

可以发现的是：

• $f[i][0]=(\sum _{k=0}^{n-1}f[i-1][k])-f[i-1][0]=(n-1)\times f[i-1][1]$
• $f[i][j]=(\sum _{k=0}^{n-1}f[i-1][k])-f[i-1][j]=(n-2)\times f[i-1][1]+f[i-1][0]$，其中 $1\le j<n$

这意味着我们其实对于每一层只需要求两个值 $f[i][0],f[i][1]$ 即可。

即：

• $f[i][0]=(n-1)\times f[i-1][1]$
• $f[i][1]=(n-2)\times f[i-1][1]+f[i-1][0]$

但是这么做的时间复杂度为 $O(m)$，考虑如何优化。

矩阵快速幂

这里令 $p_i=f[i][0],q_i=f[i][1]$

$$\left[\begin{array}{c}{k}_1{k}_2\\ k_3{k}_4\end{array}\right]\left[\begin{array}{c}{p}_{i-1}\\ q_{i-1}\end{array}\right]=\left[\begin{array}{c}{k}_1{p}_{i-1}+k_2{q}_{i-1}\\ k_3{p}_{i-1}+k_4{q}_{i-1}\end{array}\right]=\left[\begin{array}{c}{p}_i\\ q_i\end{array}\right]$$

这里定义了一个矩阵$\left[\begin{array}{c}{k}_1{k}_2\\ k_3{k}_4\end{array}\right]$

然后矩阵乘法将 $i-1$ 状态转移到 $i$ 状态。

上述我们已经定义了

• $p_i=(n-1)\times q_{i-1}$ ，故 $k_1=0,k_2=n-1$
• $q_i=(n-2)\times q_{i-1}+p_{i-1}$ ，故 $k_3=1,k_4=n-2$

实际的矩阵乘法形式为：

$$\left[\begin{array}{c}0n-1\\ 1n-2\end{array}\right]\left[\begin{array}{c}{p}_{i-1}\\ q_{i-1}\end{array}\right]=\left[\begin{array}{c}{p}_i\\ q_i\end{array}\right]$$

$$\left[\begin{array}{c}{p}_m\\ q_m\end{array}\right]=\left[\begin{array}{c}0n-1\\ 1n-2\end{array}\right]\left[\begin{array}{c}{p}_{m-1}\\ q_{m-1}\end{array}\right]={\left[\begin{array}{c}0n-1\\ 1n-2\end{array}\right]}^2\left[\begin{array}{c}{p}_{m-2}\\ q_{m-2}\end{array}\right]=\cdots ={\left[\begin{array}{c}0n-1\\ 1n-2\end{array}\right]}^m\left[\begin{array}{c}{p}_0\\ q_0\end{array}\right](p_0=1,q_0=0)$$

矩阵快速幂来计算 $matri{x}^m$ 即可。

KMP

最后，我们要统计哪些 $idx$ 满足 $ss[idx:idx+n]=t$ ，这部分用 KMP 统计即可。

时间复杂度：$O(n+8\log m)$ ，KMP 算法和矩阵快速幂的复杂度。

代码

class Solution {
public:

    static constexpr int MOD = 1e9 + 7;

    vector<vector<int>> mul(const vector<vector<int>>& A, const vector<vector<int>>& B) {
        int row = A.size(), col = B.size(), mid = A[0].size();
        vector<vector<int>> res = {{0, 0}, {0, 0}};
        for (int c = 0; c < mid; ++c)
            for (int i = 0; i < row; ++i)
                for (int j = 0; j < col; ++j)
                    res[i][j] = (1ll * A[i][c] * B[c][j] % MOD + res[i][j]) % MOD;
        return res;
    }

    vector<vector<int>> qp(vector<vector<int>>& A, long long b) {
        vector<vector<int>> res = {
                {1, 0},
                {0, 1}
        };

        while (b > 0) {
            if (b & 1) res = mul(res, A);
            A = mul(A, A);
            b >>= 1;
        }

        return res;
    }

    pair<int, int> KMP(const string& a, const string& b) {
        // 末尾添加空格，就不用判断越界了
        string s = " " + a + a + " ";
        string p = " " + b + " ";
        int n = b.size();
        int m = a.size() * 2;
        
        // 寻找 p[1,n] 在 s[1,2m] 中存在的所有索引
        // res.first 表示 s[1,n]=p[1,n] 的数量，要么为0，要么为1
        // res.second 表示 s[j:j+n]=p[1,n] 的数量，1
        pair<int, int> res = {0, 0};
        vector<int> ne(n + 10);
        for (int i = 2, j = 0; i <= n; ++i) {
            while (j > 0 && p[j + 1] != p[i]) j = ne[j];
            if (p[i] == p[j + 1]) j += 1;
            ne[i] = j;
        }
        
        // 这里细节是 i < m，这样就不会让 s[n+1,2n] 被计算
        for (int i = 1, j = 0; i < m; ++i) {
            while (j > 0 && s[i] != p[j + 1]) j = ne[j];
            if (s[i] == p[j + 1]) j += 1;
            
            if (j == n) {
                if (i == n) res.first += 1;
                else res.second += 1;
            }
        }
        
        
        return res;
    }

    int numberOfWays(string s, string t, long long k) {
        int n = s.size();

        vector<vector<int>> matrix = {
                {0, n - 1},
                {1, n - 2}
        };

        int p0 = 1, q0 = 0;
        auto res = qp(matrix, k);
        int pk = (1ll * p0 * res[0][0] % MOD + 1ll * q0 * res[0][1] % MOD) % MOD;
        int qk = (1ll * p0 * res[1][0] % MOD + 1ll * q0 * res[1][1] % MOD) % MOD;

        // KMP 计算满足条件的 idx 的数量
        auto cnt = KMP(s, t);
        int ans = (1ll * cnt.first * pk % MOD + 1ll * cnt.second * qk % MOD) % MOD;
        return ans;
    }
};