SRM 230 DIV1 Hard MultiReplacer

問題 Editorial

問題

(f: string -> string)は、全ての出現するaをarepに置き換え、bをbrepに置き換え、cをcrepに置き換える(同時に)。
(length (f^iter("a")) mod m)を求めよ。ただし(f^n)はfをn回適用する関数とする。

• arep, brep, crepはa,b,cからなる1から50文字の文字列
• 1 <= iter <= 2*10^9
• 2 <= m <= 2^31-1

解答

行列累乗するだけ。
ベクトルは{aの文字数, bの文字数, cの文字数}(それぞれmod m)

コード

#include <string>
#include <vector>
#include <algorithm>
#define rep(i,n) for(int (i)=0;(i)<(int)(n);++(i))
#define all(o) (o).begin(), (o).end()
using namespace std;
typedef long long ll;
typedef vector<ll> Vec; typedef vector<Vec> Mat;

ll MOD;

Mat identityMat(int n) {
	Mat A(n, Vec(n));
	rep(i, n) A[i][i] = 1;
	return A;
}
Mat multiply(const Mat& A, const Mat& B) {
	int n = A.size(), m = B[0].size(), p = B.size();
	Mat C(n, Vec(m));
	rep(i, n) rep(j, m) rep(k, p) {
		(C[i][j] += A[i][k] * B[k][j]) %= MOD;
	}
	return C;
}
Mat power(const Mat& A, ll n) {
	return n == 0 ? identityMat(A.size()) :
		(n & 1) ? multiply(A, power(A, n-1)) : power(multiply(A, A), n>>1);
}

struct MultiReplacer {
	long long length(string arep, string brep, string crep, int iter, int m) {
		MOD = m;
		Mat A(3, Vec(3, 0));
		string reps[3] = {arep, brep, crep};
		rep(i, 3) rep(j, 3) A[i][j] = count(all(reps[i]), 'a'+j);
		Mat V(1, Vec(3, 0));
		V[0][0] = 1;
		V = multiply(V, power(A, iter));
		return (V[0][0] + V[0][1] + V[0][2]) % MOD;
	}
};