verify/library_checker_pow_of_matrix.cpp

View this file on GitHub
Last update: 2026-05-30 18:08:28+09:00
Problem: https://judge.yosupo.jp/problem/pow_of_matrix

Depends on

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/pow_of_matrix
#include <bits/stdc++.h>
using namespace std;

#include <bits/stdc++.h>
using namespace std;
#include "templates/alias.hpp"
#include "templates/macro.hpp"

#include "matrix/matrix.hpp"

inline ll op(ll a, ll b) { return (a * b) % 998244353; }
inline ll add(ll a, ll b) { return (a + b) % 998244353; }
inline ll e() { return 0; }
inline ll id() { return 1; }

int main() {
  ll N, K;
  in(N, K);
  vvi A(N, vi(N));
  in(A);

  Matrix<ll, add, op, e, id> mat(A);
  mat ^= K;

  rep(i, N) rep(k, N) cout << mat[i][k] << (k + 1 == N ? "\n" : " ");
}

#line 1 "verify/library_checker_pow_of_matrix.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/pow_of_matrix
#include <bits/stdc++.h>
using namespace std;

#line 6 "verify/library_checker_pow_of_matrix.cpp"
using namespace std;
#line 7 "templates/alias.hpp"

template <class T> using VC = std::vector<T>;
template <class T>
using rpriority_queue = std::priority_queue<T, VC<T>, std::greater<T>>;
using ll = long long;
using ld = long double;
using pii = std::pair<ll, ll>;
using vi = VC<ll>;
using vvi = VC<vi>;
using vvvi = VC<vvi>;
using vb = VC<bool>;
using vvb = VC<vb>;
using vf = VC<double>;
using vvf = VC<vf>;
using vpii = VC<pii>;
using vvpii = VC<vpii>;
using si = std::set<ll>;
using spii = std::set<pii>;
using mii = std::map<ll, ll>;
const std::string upperlist = "ABCDEFGHIJKLMNOPQRSTUVWXYZ";
const std::string lowerlist = "abcdefghijklmnopqrstuvwxyz";
#define mp make_pair
#define dms << " " <<

constexpr int MOD998 = 998244353;
#line 4 "templates/macro.hpp"

// 引数の長さで内容が変わるrep 参考: https://trap.jp/post/1224
#define overload4(a, b, c, d, ...) d
#define _rep(i, n) for (int i = 0; i < (int)(n); i++)
#define REP(i, a, b) for (int i = (int)(a); i < (int)(b); ++i)
#define rep(...) overload4(__VA_ARGS__, REP, _rep)(__VA_ARGS__)
#define _rrep(i, n) for (int i = n - 1; i >= 0; i--)
#define RREP(i, a, b) for (int i = (int)(b - 1); i >= (int)(a); i--)
#define rrep(...) overload4(__VA_ARGS__, RREP, _rrep)(__VA_ARGS__)

#define all(a) (a).begin(), (a).end()

template <typename T> bool chmin(T &a, T b) {
  if (a > b) {
    a = b;
    return true;
  }
  return false;
}
template <typename T> bool chmax(T &a, T b) {
  if (a < b) {
    a = b;
    return true;
  }
  return false;
}

template <typename T1, typename T2>
std::istream &operator>>(std::istream &is, std::pair<T1, T2> &p) {
  is >> p.first >> p.second;
  return is;
}

template <typename T>
std::istream &operator>>(std::istream &is, std::vector<T> &v) {
  for (T &in : v)
    is >> in;
  return is;
}

template <typename T>
std::ostream &operator<<(std::ostream &os, const std::vector<T> &v) {
  for (int i = 0; i < static_cast<int>(v.size()); i++) {
    os << v[i] << (i + 1 == v.size() ? "" : " ");
  }
  return os;
}

// pythonのprintライクな関数 参考:
// https://nyaannyaan.github.io/library/template/inout.hpp
inline void out() { std::cout << std::endl; }
template <typename T, typename... U, char sep = ' '>
void out(const T &t, const U &...u) {
  std::cout << t;

  if (sizeof...(u))
    std::cout << sep;

  out(u...);
}

// cinの短縮関数 参考: https://nyaannyaan.github.io/library/template/inout.hpp
inline void in() {}
template <typename T, class... U> void in(T &t, U &...u) {
  std::cin >> t;
  in(u...);
}

template <typename T> inline T ceil_div(T a, T b) { return (a + b - 1) / b; }
template <typename T> inline T mod_pow(T a, T n, T mod) {
  T res = 1;
  while (n) {
    if (n % 2 != 0) {
      res *= a;
      res %= mod;
    }

    a *= a;
    a %= mod;

    n >>= 1;
  }

  return res;
}
template <typename T> inline T minus_mod(T a, T b) { return ((a % b) + b) % b; }

template <typename T> void apply_vec(std::vector<T> &v, T (*fn)(T)) {
  for (int i = 0; i < v.size(); i++)
    v[i] = fn(v[i]);
}
#line 9 "verify/library_checker_pow_of_matrix.cpp"

#line 3 "matrix/matrix.hpp"

template <class T, auto add, auto op, auto e, auto id> struct Matrix {
  static_assert(std::is_invocable_r_v<T, decltype(op), T, T>,
                "op must work as T(T, T)");

  static_assert(std::is_invocable_r_v<T, decltype(add), T, T>,
                "add must work as T(T, T)");

  static_assert(std::is_invocable_r_v<T, decltype(e)>, "e must work as T()");

  static_assert(std::is_invocable_r_v<T, decltype(id)>, "id must work as T()");

  std::vector<std::vector<T>> A;

  Matrix() = default;
  Matrix(int n) : A(n, std::vector<T>(n, e())) {}
  Matrix(int n, int m) : A(n, std::vector<T>(m, e())) {}
  Matrix(std::vector<std::vector<T>> A) : A(A) {}

  const inline std::vector<T> &operator[](int k) const { return A[k]; };

  inline std::vector<T> &operator[](int k) { return A[k]; };

  static Matrix<T, add, op, e, id> I(int n) {
    Matrix<T, add, op, e, id> mat(n);
    for (int i = 0; i < n; i++)
      mat[i][i] = id();

    return (mat);
  }

  Matrix &operator*=(const Matrix<T, add, op, e, id> &B) {
    int n = A.size();
    int m = B.A[0].size();
    int p = A[0].size();

    assert(n == B.A.size());
    std::vector<std::vector<T>> C(n, std::vector<T>(m, e()));

    for (int i = 0; i < n; i++) {
      for (int k = 0; k < p; k++) {
        for (int j = 0; j < m; j++) {
          C[i][j] = add(C[i][j], op((*this)[i][k], B[k][j]));
        }
      }
    }

    A.swap(C);

    return (*this);
  }

  Matrix &operator^=(long long k) {
    auto B = Matrix<T, add, op, e, id>::I(A.size());

    while (k > 0) {
      if (k & 1) {
        B *= *this;
      }

      *this *= *this;
      k >>= 1LL;
    }

    A.swap(B.A);

    return (*this);
  }

  Matrix operator*(const Matrix &B) const { return (Matrix(*this) *= B); }
  Matrix operator^(const long long k) const { return (Matrix(*this) ^= k); }
};
#line 11 "verify/library_checker_pow_of_matrix.cpp"

inline ll op(ll a, ll b) { return (a * b) % 998244353; }
inline ll add(ll a, ll b) { return (a + b) % 998244353; }
inline ll e() { return 0; }
inline ll id() { return 1; }

int main() {
  ll N, K;
  in(N, K);
  vvi A(N, vi(N));
  in(A);

  Matrix<ll, add, op, e, id> mat(A);
  mat ^= K;

  rep(i, N) rep(k, N) cout << mat[i][k] << (k + 1 == N ? "\n" : " ");
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	example_00	AC	2 ms	3 MB
g++	example_01	AC	2 ms	3 MB
g++	example_02	AC	2 ms	3 MB
g++	frobenius_hack_00	AC	2 ms	3 MB
g++	lowrank_max_random_00	AC	1826 ms	4 MB
g++	lowrank_max_random_01	AC	1671 ms	4 MB
g++	max_random_00	AC	1839 ms	4 MB
g++	max_random_01	AC	1818 ms	4 MB
g++	max_random_02	AC	1941 ms	4 MB
g++	max_random_03	AC	1817 ms	4 MB
g++	max_random_worst_00	AC	2430 ms	4 MB
g++	max_random_worst_01	AC	2432 ms	4 MB
g++	max_random_worst_02	AC	2432 ms	4 MB
g++	max_random_worst_03	AC	2432 ms	4 MB
g++	nontrivial_frobenius_form_00	AC	1768 ms	4 MB
g++	nontrivial_frobenius_form_01	AC	1859 ms	4 MB
g++	nontrivial_frobenius_form_02	AC	1918 ms	4 MB
g++	nontrivial_frobenius_form_03	AC	1816 ms	4 MB
g++	nontrivial_frobenius_form_04	AC	1676 ms	4 MB
g++	nontrivial_frobenius_form_05	AC	1777 ms	4 MB
g++	nontrivial_frobenius_form_06	AC	1839 ms	4 MB
g++	nontrivial_frobenius_form_07	AC	1718 ms	4 MB
g++	nontrivial_frobenius_form_08	AC	1940 ms	4 MB
g++	nontrivial_frobenius_form_09	AC	1898 ms	4 MB
g++	perm_max_random_00	AC	1764 ms	4 MB
g++	perm_max_random_01	AC	1546 ms	4 MB
g++	random_00	AC	1457 ms	4 MB
g++	random_01	AC	1744 ms	4 MB
g++	random_02	AC	165 ms	4 MB
g++	signed_overflow_00	AC	17 ms	4 MB
g++	small_00	AC	2 ms	3 MB
g++	small_01	AC	2 ms	3 MB
g++	small_02	AC	2 ms	3 MB
g++	small_03	AC	2 ms	3 MB
g++	small_04	AC	2 ms	3 MB
g++	small_05	AC	2 ms	3 MB
g++	small_06	AC	2 ms	3 MB
g++	small_07	AC	2 ms	3 MB
g++	small_08	AC	2 ms	3 MB
g++	small_09	AC	2 ms	3 MB
g++	small_10	AC	2 ms	3 MB
g++	small_11	AC	2 ms	3 MB
g++	unsigned_overflow_00	AC	4 ms	3 MB