본문으로 건너뛰기

3337 - Total Characters in String After Transformations II

info

풀이 키워드

스포주의

수학


풀이 코드

info
  • 메모리: 45410 KB
  • 시간: 72 ms
import java.util.*;

class Solution {
static final int MOD = 1_000_000_007;

// matrix multiplication
private long[][] matMul(long[][] a, long[][] b) {
long[][] res = new long[26][26];
for (int i = 0; i < 26; ++i) {
for (int k = 0; k < 26; ++k) {
if (a[i][k] == 0) continue;
for (int j = 0; j < 26; ++j) {
res[i][j] = (res[i][j] + a[i][k] * b[k][j]) % MOD;
}
}
}
return res;
}

// matrix exponentiation
private long[][] matPow(long[][] matrix, int power) {
long[][] res = new long[26][26];
for (int i = 0; i < 26; ++i) {
res[i][i] = 1; // init w/ identity matrix
}
while (power > 0) {
if ((power & 1) == 1) {
res = matMul(res, matrix);
}
matrix = matMul(matrix, matrix);
power >>= 1;
}
return res;
}

public int lengthAfterTransformations(String s, int t, List<Integer> nums) {
// 1. init transition matrix
long[][] T = new long[26][26];
for (int ch = 0; ch < 26; ++ch) {
int num = nums.get(ch);
for (int j = 1; j <= num; ++j) {
int nxtCh = (ch + j) % 26;
++T[ch][nxtCh];
}
}

// 2. init vector
long[] S = new long[26];
for (char ch : s.toCharArray()) ++S[ch-'a'];

long[][] T_t = matPow(T, t); // calc T^t

// calc S * T^t
long[] S_t = new long[26];
for (int i = 0; i < 26; ++i)
for (int j = 0; j < 26; ++j)
S_t[j] = (S_t[j] + S[i] * T_t[i][j]) % MOD;

// sum all vector elements
long ans = 0;
for (long v : S_t) ans = (ans + v) % MOD;

return (int)ans;
}
}

풀이 해설

transformation당 한 문자에 최대 25글자로 확장될 수 있으므로

3335번 문제처럼 DP를 통한 시뮬레이션은 O(25t)O(25^t)라서 어렵다.

대신 문자의 변환 관계가 선형 변환이기 때문에 행렬로 모델링할 수 있고,

이를 이용해 matrix exponentiation을 적용해서 풀어야 한다.


🔄 1. 문자열을 상태 벡터로 표현하기

S=[s0s1s25],si=count of letter i\mathbf{S} = \begin{bmatrix} s_0 \\ s_1 \\ \vdots \\ s_{25} \end{bmatrix},\quad s_i = \text{count of letter } i

long[] S = new long[26];
for (char ch : s.toCharArray()) ++S[ch-'a'];

🔄 2. 문자별 변환 관계를 전이 행렬로 표현하기

Tj,i={1if letter i transforms to letter j0otherwiseT_{j, i} = \begin{cases} 1 & \text{if letter } i \text{ transforms to letter } j \\ 0 & \text{otherwise} \end{cases}

예를 들어 nums[i] = 2 라면, Ti+1,i=1,  Ti+2,i=1T_{i+1, i} = 1,\;T_{i+2, i} = 1 이다.

long[][] T = new long[26][26];
for (int ch = 0; ch < 26; ++ch) {
int num = nums.get(ch);
for (int j = 1; j <= num; ++j) {
int nxtCh = (ch + j) % 26;
++T[ch][nxtCh];
}
}

결론적으로 답은 1T(TtS)\mathbf{1}^T \cdot ( \mathbf{T}^t \cdot \mathbf{S} ) 이고,

naive하게 TtT^t 계산 기준으로 따지자면 시간복잡도는 O(logt)O(log\,t)가 된다.


메모

  • 코테에 나올리가 없다. 잘 가다가 왜 또 문제가 급발진하는?
  • 모듈로 값을 static으로 선언하면 조오금 더 빨라져요