【模板】矩阵加、乘、快速幂 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;
}