POJ 3735 Training little cats

問題文

DP強化週間。

猫がN匹いる。3種類の命令が存在し、それが「g i」ならi番目の猫の持つピーナツをひとつ増やす。「e i」ならi番目の猫が持つピーナツの数をゼロにする。「s i j」ならi番目とj番目の猫が所有するピーナツを交換する。これらの命令が１セット中にk回与えられる。ｍセット行った後に各猫が持つピーナツの数を求める問題。

愚直にシミュレートすることは制限時間上できない。
1セットごとに状態が遷移するので漸化式を考える。すると漸化式の係数をまとめたものはN個の漸化式分と定数項を一つ加えた(N+1)×(N+1)の行列となる。あとはこの行列のM乗を計算して定数項のところをとり出せばいい。

という漸化式に最初なっているとして、たとえば「s 1 2」は行列の1列目と2列目を入れ替えて

にすることに相当する。

普通に書くとTLE。掛け算のところを少し工夫すると速くなって通った。

typedef long long ll;

int N;
ll A[101][101], B[101][101], C[101][101];

void mat_mul(ll a[][101], ll b[][101], ll dst[][101]){
  int i, j, k;
  memset(C, 0, sizeof(C));
  for(i = 0; i < N + 1; i++)
    for(k = 0; k < N + 1; k++)
      if(a[i][k])//This makes Mat multiplication much faster
      for(j = 0; j < N + 1; j++)
	C[i][j] += a[i][k] * b[k][j];

  for(i = 0; i < N + 1; i++)
    for(j = 0; j < N + 1; j++)
      dst[i][j] = C[i][j];
}

void mat_pow(ll A[][101], ll n){
  int i;
  memset(B, 0, sizeof(B));
  for(i = 0; i < N + 1; i++){
    B[i][i] = 1;
  }
  while(n > 0){
    if(n & 1) mat_mul(B, A, B);
    mat_mul(A, A, A);
    n >>= 1;
  }
}

int main(int argc, char **argv){
  ll n, m, k;
  int i, j;
  while(scanf("%lld %lld %lld", &n, &m, &k), n|m|k){
    N = n;
    memset(A, 0, sizeof(A));
    for(i = 0; i < n + 1; i++) A[i][i] = 1;
    for(i = 0; i < k; i++){
      char a[5];
      int b, c;
      scanf("%s", &a);
      if(a[0] == 'g'){
	scanf("%d", &b);
	A[n][b - 1]++;
      }else if(a[0] == 'e'){
	scanf("%d", &b);
	for(j = 0; j < n + 1; j++) 
	  A[j][b - 1] = 0;
      }else{
	scanf("%d %d", &b, &c);
	for(j = 0; j < n + 1; j++) {
	  ll tmp = A[j][b - 1];
	  A[j][b - 1] = A[j][c - 1];
	  A[j][c - 1] = tmp;
	}
      }
    }
    mat_pow(A, m);
    for(i = 0; i < n; i++){
      printf("%lld%c", B[n][i], (i == n - 1) ? '\n':' ');
    }
  }
}

強化週間は引き続きやっていこうと思う。