test/math/matrix.test.cpp

• View this file on GitHub
• Last update: 2025-11-15 20:52:05+09:00
• Problem: https://judge.yosupo.jp/problem/pow_of_matrix

Depends on

• Matrix (行列) (math/matrix.hpp)
• Static Mod Int (math/static_mod_int.hpp)
• template/constant.hpp
• template/include.hpp
• template/macro.hpp
• 提出用テンプレート (template/template.hpp)
• template/type_alias.hpp
• Choose Minimum / Maximum (utility/choose_min_max.hpp)
• utility/set_io.hpp

Required by

• test/utility/next_combination.test.cpp

Code

#define PROBLEM "https://judge.yosupo.jp/problem/pow_of_matrix"

#include "../../math/matrix.hpp"
#include "../../math/static_mod_int.hpp"
#include "../../template/template.hpp"

using mint = StaticModInt<998244353>;

int main() {
    set_io();

    int N;
    cin >> N;

    ll K;
    cin >> K;

    Matrix<mint> mat(N);
    rep (i, 0, N) {
        rep (j, 0, N) {
            ll a;
            cin >> a;

            mat.set(i, j, mint(a));
        }
    }

    mat = mat.pow(K);

    rep (i, 0, N) {
        rep (j, 0, N) {
            if (j > 0) {
                cout << ' ';
            }

            cout << mat.get(i, j).val();
        }

        cout << LF;
    }

    return 0;
}

#line 1 "test/math/matrix.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/pow_of_matrix"

#line 2 "math/matrix.hpp"

#line 2 "template/template.hpp"

/**
 * @author ku_senjan
 * @title 提出用テンプレート
 * @see https://github.com/kogetsu0728/ku-library
 */

#line 2 "template/constant.hpp"

#line 2 "template/include.hpp"

#ifdef LOCAL
#define _GLIBCXX_DEBUG
#endif

#include <bits/stdc++.h>
using namespace std;
#line 2 "template/type_alias.hpp"

#line 4 "template/type_alias.hpp"

using uint = unsigned int;
using ll = long long;
using ull = unsigned long long;

using ld = long double;

template <class T, bool REVERSE = false>
using heap = priority_queue<T, vector<T>, conditional_t<REVERSE, greater<T>, less<T>>>;
#line 5 "template/constant.hpp"

template <class T>
inline constexpr T INF = numeric_limits<T>::max() / 2;

inline constexpr array<int, 4> DY4 = {0, -1, 0, 1};
inline constexpr array<int, 4> DX4 = {1, 0, -1, 0};
inline constexpr array<int, 8> DY8 = {0, -1, -1, -1, 0, 1, 1, 1};
inline constexpr array<int, 8> DX8 = {1, 1, 0, -1, -1, -1, 0, 1};

inline constexpr char LF = '\n';
#line 2 "template/macro.hpp"

#line 5 "template/macro.hpp"

/**
 * @see https://trap.jp/post/1224/
 */

#ifdef LOCAL
inline constexpr bool IS_LOCAL = true;
#else
inline constexpr bool IS_LOCAL = false;
#endif

#define IF_LOCAL if constexpr (IS_LOCAL)

#define all(a) begin(a), end(a)
#define rall(a) rbegin(a), rend(a)

#define overload4(a, b, c, d, e, ...) e

#define rep1(i, a) for (ll i = 0; (i) < (ll)(a); ++(i))
#define rep2(i, a, b) for (ll i = (ll)(a); (i) < (ll)(b); ++(i))
#define rep3(i, a, b, c) for (ll i = (ll)(a); (i) < (ll)(b); (i) += (ll)(c))
#define rep(...) overload4(__VA_ARGS__, rep3, rep2, rep1)(__VA_ARGS__)

#define rrep1(i, a) for (ll i = (ll)(a); (i) >= 0; --(i))
#define rrep2(i, a, b) for (ll i = (ll)(a); (i) >= (ll)(b); --(i))
#define rrep3(i, a, b, c) for (ll i = (ll)(a); (i) >= (ll)(b); (i) -= (ll)(c))
#define rrep(...) overload4(__VA_ARGS__, rrep3, rrep2, rrep1)(__VA_ARGS__)
#line 13 "template/template.hpp"

#line 2 "utility/choose_min_max.hpp"

/**
 * @title Choose Minimum / Maximum
 */

template <class T>
bool chmin(T& a, const T& b) {
    if (a > b) {
        a = b;
        return true;
    }

    return false;
}

template <class T>
bool chmax(T& a, const T& b) {
    if (a < b) {
        a = b;
        return true;
    }

    return false;
}
#line 2 "utility/set_io.hpp"

#line 4 "utility/set_io.hpp"

void set_io(int d = 16) {
    cin.tie(nullptr);
    ios_base::sync_with_stdio(false);
    cout << fixed << setprecision(d);

    return;
}
#line 4 "math/matrix.hpp"

/**
 * @brief Matrix (行列)
 */
template <typename T>
class Matrix {
  private:
    vector<vector<T>> data;

  public:
    Matrix() : Matrix(0) {}
    explicit Matrix(int _h) : Matrix(_h, _h) {}
    explicit Matrix(int _h, int _w) : data(_h, vector<T>(_w, T{})) {}

    //! 単位行列
    static Matrix identity(int s) {
        Matrix res(s);

        rep (i, 0, s) {
            res.set(i, i, T(1));
        }

        return res;
    }

    int row() const {
        return data.size();
    }

    int col() const {
        return data.empty() ? 0 : data[0].size();
    }

    T get(int i, int j) const {
        assert(0 <= i && i < row());
        assert(0 <= j && j < col());

        return data[i][j];
    }

    void set(int i, int j, const T v) {
        assert(0 <= i && i < row());
        assert(0 <= j && j < col());

        data[i][j] = v;

        return;
    }

    friend bool operator==(const Matrix& lhs, const Matrix& rhs) {
        return lhs.data == rhs.data;
    }

    friend bool operator!=(const Matrix& lhs, const Matrix& rhs) {
        return lhs.data != rhs.data;
    }

    friend Matrix operator*(const Matrix& lhs, const Matrix& rhs) {
        assert(lhs.col() == rhs.row());

        Matrix res(lhs.row(), rhs.col());
        rep (i, 0, lhs.row()) {
            rep (j, 0, rhs.col()) {
                rep (k, 0, lhs.col()) {
                    res.set(i, j, res.get(i, j) + lhs.get(i, k) * rhs.get(k, j));
                }
            }
        }

        return res;
    }

    Matrix& operator*=(const Matrix& rhs) {
        return *this = *this * rhs;
    }

    Matrix pow(ll y) const {
        assert(row() == col());
        assert(0 <= y);

        Matrix res = identity(row());
        Matrix x = *this;

        while (y > 0) {
            if ((y & 1) != 0) {
                res *= x;
            }

            x *= x;
            y >>= 1;
        }

        return res;
    }
};
#line 2 "math/static_mod_int.hpp"

#line 4 "math/static_mod_int.hpp"

/**
 * @brief Static Mod Int
 */
template <int M>
class StaticModInt {
    static_assert(2 < M);
    static_assert(M <= INT_MAX / 2);

  private:
    int v;

  public:
    StaticModInt() : v(0) {}
    StaticModInt(ll _v) : v(int(((_v % M) + M) % M)) {}

    static StaticModInt raw(int _v) {
        StaticModInt res;
        res.v = _v;

        return res;
    }

    static int mod() { return M; }
    int val() const { return v; }

    friend bool operator==(const StaticModInt& lhs, const StaticModInt& rhs) {
        return lhs.v == rhs.v;
    }

    friend bool operator!=(const StaticModInt& lhs, const StaticModInt& rhs) {
        return !(lhs == rhs);
    }

    StaticModInt& operator+=(const StaticModInt& rhs) {
        v += rhs.val();
        if (M <= v) {
            v -= M;
        }

        return *this;
    }

    StaticModInt& operator-=(const StaticModInt& rhs) {
        if (v < rhs.val()) {
            v += M;
        }
        v -= rhs.val();

        return *this;
    }

    StaticModInt& operator*=(const StaticModInt& rhs) {
        v = (ull(v) * rhs.val()) % M;

        return *this;
    }

    StaticModInt& operator/=(const StaticModInt& rhs) {
        assert(rhs.val() != 0);

        return *this *= rhs.inv();
    }

    friend StaticModInt operator+(const StaticModInt& lhs, const StaticModInt& rhs) {
        return StaticModInt(lhs) += rhs;
    }

    friend StaticModInt operator-(const StaticModInt& lhs, const StaticModInt& rhs) {
        return StaticModInt(lhs) -= rhs;
    }

    friend StaticModInt operator*(const StaticModInt& lhs, const StaticModInt& rhs) {
        return StaticModInt(lhs) *= rhs;
    }

    friend StaticModInt operator/(const StaticModInt& lhs, const StaticModInt& rhs) {
        return StaticModInt(lhs) /= rhs;
    }

    StaticModInt pow(ll y) const {
        StaticModInt res = raw(1);
        StaticModInt x = *this;

        while (0 < y) {
            if (y & 1) {
                res *= x;
            }

            x *= x;
            y >>= 1;
        }

        return res;
    }

    StaticModInt inv() const { return pow(M - 2); }
};
#line 6 "test/math/matrix.test.cpp"

using mint = StaticModInt<998244353>;

int main() {
    set_io();

    int N;
    cin >> N;

    ll K;
    cin >> K;

    Matrix<mint> mat(N);
    rep (i, 0, N) {
        rep (j, 0, N) {
            ll a;
            cin >> a;

            mat.set(i, j, mint(a));
        }
    }

    mat = mat.pow(K);

    rep (i, 0, N) {
        rep (j, 0, N) {
            if (j > 0) {
                cout << ' ';
            }

            cout << mat.get(i, j).val();
        }

        cout << LF;
    }

    return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	example_00	AC	6 ms	4 MB
g++	example_01	AC	5 ms	4 MB
g++	example_02	AC	5 ms	4 MB
g++	frobenius_hack_00	AC	5 ms	4 MB
g++	lowrank_max_random_00	AC	2413 ms	4 MB
g++	lowrank_max_random_01	AC	2199 ms	4 MB
g++	max_random_00	AC	2414 ms	4 MB
g++	max_random_01	AC	2390 ms	4 MB
g++	max_random_02	AC	2555 ms	4 MB
g++	max_random_03	AC	2395 ms	4 MB
g++	max_random_worst_00	AC	3199 ms	4 MB
g++	max_random_worst_01	AC	3193 ms	4 MB
g++	max_random_worst_02	AC	3191 ms	4 MB
g++	max_random_worst_03	AC	3200 ms	4 MB
g++	nontrivial_frobenius_form_00	AC	2334 ms	4 MB
g++	nontrivial_frobenius_form_01	AC	2441 ms	4 MB
g++	nontrivial_frobenius_form_02	AC	2526 ms	4 MB
g++	nontrivial_frobenius_form_03	AC	2400 ms	4 MB
g++	nontrivial_frobenius_form_04	AC	2209 ms	4 MB
g++	nontrivial_frobenius_form_05	AC	2302 ms	4 MB
g++	nontrivial_frobenius_form_06	AC	2418 ms	4 MB
g++	nontrivial_frobenius_form_07	AC	2251 ms	4 MB
g++	nontrivial_frobenius_form_08	AC	2548 ms	4 MB
g++	nontrivial_frobenius_form_09	AC	2502 ms	4 MB
g++	perm_max_random_00	AC	2324 ms	4 MB
g++	perm_max_random_01	AC	2043 ms	4 MB
g++	random_00	AC	1915 ms	4 MB
g++	random_01	AC	2298 ms	4 MB
g++	random_02	AC	219 ms	4 MB
g++	signed_overflow_00	AC	24 ms	4 MB
g++	small_00	AC	4 ms	4 MB
g++	small_01	AC	5 ms	4 MB
g++	small_02	AC	4 ms	4 MB
g++	small_03	AC	4 ms	4 MB
g++	small_04	AC	4 ms	4 MB
g++	small_05	AC	4 ms	4 MB
g++	small_06	AC	4 ms	4 MB
g++	small_07	AC	4 ms	4 MB
g++	small_08	AC	4 ms	4 MB
g++	small_09	AC	4 ms	3 MB
g++	small_10	AC	4 ms	4 MB
g++	small_11	AC	4 ms	4 MB
g++	unsigned_overflow_00	AC	6 ms	4 MB

Back to top page