String Transformation

You are given two strings s and t of equal length n. You can perform the following operation on the string s:

Remove a suffix of s of length l where 0 < l < n and append it at the start of s.

For example, let s = 'abcd' then in one operation you can remove the suffix 'cd' and append it in front of s making s = 'cdab'.

You are also given an integer k. Return the number of ways in which s can be transformed into t in exactly k operations.

Since the answer can be large, return it modulo 10⁹ + 7 .

Example 1:

Input: s = "abcd", t = "cdab", k = 2
Output: 2
Explanation: 
First way:
In first operation, choose suffix from index = 3, so resulting s = "dabc".
In second operation, choose suffix from index = 3, so resulting s = "cdab".

Second way:
In first operation, choose suffix from index = 1, so resulting s = "bcda".
In second operation, choose suffix from index = 1, so resulting s = "cdab".

Example 2:

Input: s = "ababab", t = "ababab", k = 1
Output: 2
Explanation: 
First way:
Choose suffix from index = 2, so resulting s = "ababab".

Second way:
Choose suffix from index = 4, so resulting s = "ababab".

Constraints:

2 <= s.length <= 5 * 10⁵
1 <= k <= 10¹⁵
s.length == t.length
s and t consist of only lowercase English alphabets.

解题思路

　　比赛看到$k$的数据范围就猜到是矩阵乘法加快速幂的做法，不过还是没做出来。

　　把某个后缀移到字符串的最前面所得到的结果，本质上就是把字符串看成一个环，然后选择某个间隙把环拆开（不能选择原字符串中第一个与最后一个字符的间隙），得到一条链。假设字符串的长度为$n$，第一个字符是$s_0$，那么通过一次操作得到结果就是在环中选择一个$s_i, \, i \in [1,n-1]$作为开头，然后依次选择下一个字符$s_{(i+1) \bmod n}$拼接，展开成新的字符串。可以发现无论操作了多少次，所有结果形成的环都是一样的。

　　我们先统计原来的$s$中有多少个下标$i$满足在操作一次后，以$s_i$为开头的字符串与$t$匹配。假设为$c$，那么不满足的下标数量自然就是$n-c$了。统计的方法很简单，由于涉及到环，先在$s$后面再拼接一个$s$，然后对长度为$2n$的拼接字符串跑一遍kmp，统计在下标$0 \sim n-1$中有多少个位置与$t$匹配即可（代码中为了方便下标从$1$开始）。

　　题目问恰好操作$k$次后$s$与$t$匹配的方案数，等价于问操作了$k$次后以原串$s$中那$c$个下标为开头的方案数是多少。所以我们可以先定义状态$f(i)$表示操作了$i$次后，得到以那$c$个下标为开头的字符串的方案数。可以从操作了$i-1$次的状态转移过来，但又发现在$i-1$的状态中，既可以从以那$c$个下标为开头的字符串转移过来，也可以从另外的$n-c$个下标为开头的字符串转移过来，所以我们再定义另外一个状态$g(i)$表示操作了$i$次后，得到以那$n-c$个下标为开头的字符串的方案数。因此就有状态转移方程$$\Big\{\begin{array}{left} f(i) = f(i-1) \cdot (c-1) + g(i-1) \cdot c \\ g(i) = f(i-1) \cdot (n-c) + g(i-1) \cdot (n-c-1) \end{array}$$

　　如果直接算的话时间复杂度是$O(k)$的，为此我们试一下能不能转换成矩阵乘法。我们尝试把上面的式子变成矩阵乘法的形式：$$\begin{bmatrix} f(i) & g(i) \end{bmatrix} = \begin{bmatrix} f(i-1) & g(i-1) \end{bmatrix} \times \begin{bmatrix} c-1 & n-c \\ c & n-c-1 \end{bmatrix}$$

　　发现是可以的，因此定义$F_i = \begin{bmatrix} f(i) & g(i) \end{bmatrix}$，系数矩阵$A = \begin{bmatrix} c-1 & n-c \\ c & n-c-1 \end{bmatrix}$，那么就会有$F_i = F_{i-1} \times A$。根据矩阵乘法的结合律，那么就有$F_{k} = F_0 \times A^k $。根据定义，$F_0$表示没有任何操作时的状态，因此如果初始时$s=t$，那么就有$F_0 = \begin{bmatrix} 1 & 0 \end{bmatrix}$，否则$F_0 = \begin{bmatrix} 0 & 1 \end{bmatrix}$。答案就是$F_{k}$中的$f(k)$，用快速幂求$F_0 \times A^k $即可。

　　为了统一成矩阵乘矩阵的形式，代码中把$F_i$扩充成了$2 \times 2$的矩阵，具体计算如下：$$\begin{bmatrix} f(i) & g(i) \\ 0 & 0 \end{bmatrix} = \begin{bmatrix} f(i-1) & g(i-1) \\ 0 & 0 \end{bmatrix} \times \begin{bmatrix} c-1 & n-c \\ c & n-c-1 \end{bmatrix}$$

　　AC代码如下，时间复杂度为$O(n + \log{k})$：

 1 class Solution {
 2 public:
 3     int mod = 1e9 + 7;
 4     
 5     int numberOfWays(string s, string t, long long k) {
 6         int n = t.size();
 7         vector<int> ne(n + 1);
 8         for (int i = 2, j = 0; i <= n; i++) {
 9             while (j && t[j] != t[i - 1]) {
10                 j = ne[j];
11             }
12             if (t[j] == t[i - 1]) j++;
13             ne[i] = j;
14         }
15         string ss = s + s;
16         int c = 0;
17         for (int i = 1, j = 0; i <= n << 1; i++) {
18             while (j && t[j] != ss[i - 1]) {
19                 j = ne[j];
20             }
21             if (t[j] == ss[i - 1]) j++;
22             if (j == n) {
23                 if (i - n + 1 <= n) c++;
24                 j = ne[j];
25             }
26         }
27         vector<vector<int>> f({{s == t, s != t}, {0, 0}}), a({{c - 1, n - c}, {c, n - c - 1}});
28         function<void(vector<vector<int>>&, vector<vector<int>>&, vector<vector<int>>&)> mul = [&](vector<vector<int>> &c, vector<vector<int>> &a, vector<vector<int>> &b) {
29             vector<vector<int>> tmp(2, vector<int>(2));
30             for (int i = 0; i < 2; i++) {
31                 for (int j = 0; j < 2; j++) {
32                     for (int k = 0; k < 2; k++) {
33                         tmp[i][j] = (tmp[i][j] + 1ll * a[i][k] * b[k][j]) % mod;
34                     }
35                 }
36             }
37             c = tmp;
38         };
39         while (k) {
40             if (k & 1) mul(f, f, a);
41             mul(a, a, a);
42             k >>= 1;
43         }
44         return f[0][0];
45     }
46 };