【模板】矩阵加、乘、快速幂 matrix

posted on 2022-01-25 18:52:13 | under 模板 | source

typedef long long LL;
const int P=998244353;
template<int N,int M,class T=LL> struct matrix{
	T a[N+1][M+1];
	matrix(bool k=0){memset(a,0,sizeof a);for(int i=1;i<=N;i++) a[i][i]=k;}
	T* operator[](int i){return a[i];}
};
template<int N,int M,class T=LL> matrix<N,M,T> operator+(matrix<N,M,T> a,matrix<N,M,T> b){
	for(int i=1;i<=N;i++)
		for(int j=1;j<=M;j++)
			a[i][j]=(a[i][j]+b[i][j])%P;
	return a;
}
template<int N,int M,int R,class T=LL> matrix<N,R,T> operator*(matrix<N,M,T> a,matrix<M,R,T> b){
	matrix<N,R,T> r=0;
	for(int i=1;i<=N;i++)
		for(int j=1;j<=R;j++)
			for(int k=1;k<=M;k++)
				r[i][j]=(r[i][j]+a[i][k]*b[k][j])%P;
	return r;
}
template<int N,class T=LL> matrix<N,N,T> operator^(matrix<N,N,T> a,LL b){
	matrix<N,N,T> r=1;
	for(;b;b>>=1,a=a*a) if(b&1) r=r*a;
	return r;
}

struct matrix {
  int n, m;
  vector<mint> vec;
  auto operator[](int i) { return vec.begin() + i * m; }
  auto operator[](int i) const { return vec.begin() + i * m; }
  matrix() = default;
  matrix(int _n, int _m) : n(_n), m(_m), vec(_n * _m, 0) {}
};
matrix operator*(const matrix& lhs, const matrix& rhs) {
  assert(lhs.m == rhs.n);
  matrix ret(lhs.n, rhs.m);
  for (int i = 0; i < lhs.n; i++) {
    for (int j = 0; j < lhs.m; j++) {
      for (int k = 0; k < rhs.m; k++) ret[i][k] += lhs[i][j] * rhs[j][k];
    }
  }
  return ret;
}
matrix I(int n) {
  matrix ret(n, n);
  for (int i = 0; i < n; i++) ret[i][i] = 1;
  return ret;
}