Skip to content

3337. Total Characters in String After Transformations II 👍

  • Time: $O(|\texttt{s}| + \log t)$
  • Space: $O(26^2 + 26) = O(1)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
class Solution {
 public:
  // Similar to 3335. Total Characters in String After Transformations I
  int lengthAfterTransformations(string s, int t, vector<int>& nums) {
    // T[i][j] := the number of ways to transform ('a' + i) to ('a' + j)
    const vector<vector<int>> T = getTransformationMatrix(nums);
    const vector poweredT = matrixPow(T, t);
    vector<int> count(26);
    // lengths[i] := the total length of ('a' + i) after t transformations
    vector<long> lengths(26);

    for (const char c : s)
      ++count[c - 'a'];

    for (int i = 0; i < 26; ++i)
      for (int j = 0; j < 26; ++j) {
        lengths[j] += static_cast<long>(count[i]) * poweredT[i][j];
        lengths[j] %= kMod;
      }

    return accumulate(lengths.begin(), lengths.end(), 0L) % kMod;
  }

 private:
  static constexpr int kMod = 1'000'000'007;

  vector<vector<int>> getTransformationMatrix(const vector<int>& nums) {
    vector<vector<int>> T(26, vector<int>(26));
    for (int i = 0; i < nums.size(); ++i)
      for (int step = 1; step <= nums[i]; ++step)
        ++T[i][(i + step) % 26];
    return T;
  }

  vector<vector<int>> getIdentityMatrix(int sz) {
    vector<vector<int>> I(sz, vector<int>(sz));
    for (int i = 0; i < sz; ++i)
      I[i][i] = 1;
    return I;
  }

  // Returns A * B.
  vector<vector<int>> matrixMult(const vector<vector<int>>& A,
                                 const vector<vector<int>>& B) {
    const int sz = A.size();
    vector<vector<int>> C(sz, vector<int>(sz));
    for (int i = 0; i < sz; ++i)
      for (int j = 0; j < sz; ++j)
        for (int k = 0; k < sz; ++k)
          C[i][j] = (C[i][j] + static_cast<long>(A[i][k]) * B[k][j]) % kMod;
    return C;
  }

  // Returns M^n.
  vector<vector<int>> matrixPow(const vector<vector<int>>& M, int n) {
    if (n == 0)
      return getIdentityMatrix(M.size());
    if (n % 2 == 1)
      return matrixMult(M, matrixPow(M, n - 1));
    return matrixPow(matrixMult(M, M), n / 2);
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
class Solution {
  // Similar to 3335. Total Characters in String After Transformations I
  public int lengthAfterTransformations(String s, int t, List<Integer> nums) {
    // T[i][j] := the number of ways to transform ('a' + i) to ('a' + j)
    int[][] T = getTransformationMatrix(nums);
    int[][] poweredT = matrixPow(T, t);
    int[] count = new int[26];
    // lengths[i] := the total length of ('a' + i) after t transformations
    long[] lengths = new long[26];

    for (final char c : s.toCharArray())
      ++count[c - 'a'];

    for (int i = 0; i < 26; ++i)
      for (int j = 0; j < 26; ++j) {
        lengths[j] += (long) count[i] * poweredT[i][j];
        lengths[j] %= MOD;
      }

    return (int) (Arrays.stream(lengths).sum() % MOD);
  }

  private static final int MOD = 1_000_000_007;

  private int[][] getTransformationMatrix(List<Integer> nums) {
    int[][] T = new int[26][26];
    for (int i = 0; i < nums.size(); ++i)
      for (int step = 1; step <= nums.get(i); ++step)
        ++T[i][(i + step) % 26];
    return T;
  }

  private int[][] getIdentityMatrix(int sz) {
    int[][] I = new int[sz][sz];
    for (int i = 0; i < sz; ++i)
      I[i][i] = 1;
    return I;
  }

  // Returns A * B.
  private int[][] matrixMult(int[][] A, int[][] B) {
    final int sz = A.length;
    int[][] C = new int[sz][sz];
    for (int i = 0; i < sz; ++i)
      for (int j = 0; j < sz; ++j)
        for (int k = 0; k < sz; ++k)
          C[i][j] = (int) ((C[i][j] + (long) A[i][k] * B[k][j]) % MOD);
    return C;
  }

  // Returns M^n.
  private int[][] matrixPow(int[][] M, int n) {
    if (n == 0)
      return getIdentityMatrix(M.length);
    if (n % 2 == 1)
      return matrixMult(M, matrixPow(M, n - 1));
    return matrixPow(matrixMult(M, M), n / 2);
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
class Solution:
  # Similar to 3335. Total Characters in String After Transformations I
  def lengthAfterTransformations(self, s: str, t: int, nums: list[int]) -> int:
    MOD = 1_000_000_007

    def matrixMult(A: list[list[int]], B: list[list[int]]) -> list[list[int]]:
      """Returns A * B."""
      sz = len(A)
      C = [[0] * sz for _ in range(sz)]
      for i in range(sz):
        for j in range(sz):
          for k in range(sz):
            C[i][j] += A[i][k] * B[k][j]
            C[i][j] %= MOD
      return C

    def matrixPow(M: list[list[int]], n: int) -> list[list[int]]:
      """Returns M^n."""
      if n == 0:
        return [[1 if i == j else 0  # identity matrix
                for j in range(len(M))]
                for i in range(len(M))]
      if n % 2 == 1:
        return matrixMult(M, matrixPow(M, n - 1))
      return matrixPow(matrixMult(M, M), n // 2)

    # T[i][j] := the number of ways to transform ('a' + i) to ('a' + j)
    T = self._getTransformationMatrix(nums)
    poweredT = matrixPow(T, t)
    count = [0] * 26
    lengths = [0] * 26

    for c in s:
      count[ord(c) - ord('a')] += 1

    for i in range(26):
      for j in range(26):
        lengths[j] += count[i] * poweredT[i][j]
        lengths[j] %= MOD

    return sum(lengths) % MOD

  def _getTransformationMatrix(self, nums: list[int]) -> list[list[int]]:
    T = [[0] * 26 for _ in range(26)]
    for i, steps in enumerate(nums):
      for step in range(1, steps + 1):
        T[i][(i + step) % 26] += 1
    return T