|
| 1 | +""" |
| 2 | +DP, Z-algorithm, Fast mod. |
| 3 | +Approach |
| 4 | +How to represent a string? |
| 5 | +Each operation is just a rotation. Each result string can be represented by an integer from 0 to n - 1. Namely, it's just the new index of s[0]. |
| 6 | +How to find the integer(s) that can represent string t? |
| 7 | +Create a new string s + t + t (length = 3 * n). |
| 8 | +Use Z-algorithm (or KMP), for each n <= index < 2 * n, calculate the maximum prefix length that each substring starts from index can match, if the length >= n, then (index - n) is a valid integer representation. |
| 9 | +How to get the result? |
| 10 | +It's a very obvious DP. |
| 11 | +If we use an integer to represent a string, we only need to consider the transition from zero to non-zero and from non-zero to zero. In other words, all the non-zero strings should have the same result. |
| 12 | +So let dp[t][i = 0/1] be the number of ways to get the zero/nonzero string |
| 13 | +after excatly t steps. |
| 14 | +Then |
| 15 | +dp[t][0] = dp[t - 1][1] * (n - 1). |
| 16 | +All the non zero strings can make it. |
| 17 | +dp[t][1] = dp[t - 1][0] + dp[t - 1] * (n - 2). |
| 18 | +For a particular non zero string, all the other non zero strings and zero string can make it. |
| 19 | +We have dp[0][0] = 1 and dp[0][1] = 0 |
| 20 | +Use matrix multiplication. |
| 21 | +How to calculate dp[k][x = 0, 1] faster? |
| 22 | +Use matrix multiplication |
| 23 | +vector (dp[t - 1][0], dp[t - 1][1]) |
| 24 | +multiplies matrix |
| 25 | +[0 1] |
| 26 | +[n - 1 n - 2] |
| 27 | +== vector (dp[t][0], dp[t - 1][1]). |
| 28 | +So we just need to calculate the kth power of the matrix which can be done by fast power algorith. |
| 29 | +Complexity |
| 30 | +Time complexity: |
| 31 | +O(n + logk) |
| 32 | +Space complexity: |
| 33 | +O(n) |
| 34 | +""" |
| 35 | + |
| 36 | + |
| 37 | +class Solution: |
| 38 | + M: int = 1000000007 |
| 39 | + |
| 40 | + def add(self, x: int, y: int) -> int: |
| 41 | + x += y |
| 42 | + if x >= self.M: |
| 43 | + x -= self.M |
| 44 | + return x |
| 45 | + |
| 46 | + def mul(self, x: int, y: int) -> int: |
| 47 | + return int(x * y % self.M) |
| 48 | + |
| 49 | + def getZ(self, s: str) -> List[int]: |
| 50 | + n = len(s) |
| 51 | + z = [0] * n |
| 52 | + left = right = 0 |
| 53 | + for i in range(1, n): |
| 54 | + if i <= right and z[i - left] <= right - i: |
| 55 | + z[i] = z[i - left] |
| 56 | + else: |
| 57 | + z_i = max(0, right - i + 1) |
| 58 | + while i + z_i < n and s[i + z_i] == s[z_i]: |
| 59 | + z_i += 1 |
| 60 | + z[i] = z_i |
| 61 | + if i + z[i] - 1 > right: |
| 62 | + left = i |
| 63 | + right = i + z[i] - 1 |
| 64 | + return z |
| 65 | + |
| 66 | + def matrixMultiply(self, a: List[List[int]], b: List[List[int]]) -> List[List[int]]: |
| 67 | + m = len(a) |
| 68 | + n = len(a[0]) |
| 69 | + p = len(b[0]) |
| 70 | + r = [[0] * p for _ in range(m)] |
| 71 | + for i in range(m): |
| 72 | + for j in range(p): |
| 73 | + for k in range(n): |
| 74 | + r[i][j] = self.add(r[i][j], self.mul(a[i][k], b[k][j])) |
| 75 | + return r |
| 76 | + |
| 77 | + def matrixPower(self, a: List[List[int]], y: int) -> List[List[int]]: |
| 78 | + n = len(a) |
| 79 | + r = [[0] * n for _ in range(n)] |
| 80 | + for i in range(n): |
| 81 | + r[i][i] = 1 |
| 82 | + x = [a[i][:] for i in range(n)] |
| 83 | + while y > 0: |
| 84 | + if y & 1: |
| 85 | + r = self.matrixMultiply(r, x) |
| 86 | + x = self.matrixMultiply(x, x) |
| 87 | + y >>= 1 |
| 88 | + return r |
| 89 | + |
| 90 | + def numberOfWays(self, s: str, t: str, k: int) -> int: |
| 91 | + n = len(s) |
| 92 | + dp = self.matrixPower([[0, 1], [n - 1, n - 2]], k)[0] |
| 93 | + s += t + t |
| 94 | + z = self.getZ(s) |
| 95 | + m = n + n |
| 96 | + result = 0 |
| 97 | + for i in range(n, m): |
| 98 | + if z[i] >= n: |
| 99 | + result = self.add(result, dp[0] if i - n == 0 else dp[1]) |
| 100 | + return result |
0 commit comments