test/Math/pow_of_matrix/yosupo-pow_of_matrix.cpp

View this file on GitHub
Last update: 2025-04-13 19:22:19+09:00
Problem: https://judge.yosupo.jp/problem/pow_of_matrix

Depends on

繰り返し二乗法 (Other/fast_power.hpp)

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/pow_of_matrix

#include "Other/fast_power.hpp"
#include <bits/stdc++.h>
#include <atcoder/all>

using namespace std;

int main()
{
    long long N, K;
    cin >> N >> K;

    using S = vector<vector<atcoder::modint>>;
    auto mul = [&](S a, S b) -> S
    {
        S result(N, vector<atcoder::modint>(N));
        for (int i = 0; i < N; i++)
        {
            for (int j = 0; j < N; j++)
            {
                for (int k = 0; k < N; k++)
                {
                    result[i][j] += a[i][k] * b[k][j];
                }
            }
        }

        return result;
    };

    auto e = [&]() -> S
    {
        S result = S(N, vector<atcoder::modint>(N));
        for (int i = 0; i < N; i++)
        {
            result[i][i] = 1;
        }

        return result;
    };

    vector a(N, vector<atcoder::modint>(N));
    for (int i = 0; i < N; i++)
    {
        for (int j = 0; j < N; j++)
        {
            int in_a;
            cin >> in_a;
            a[i][j] = in_a;
        }
    }

    auto ans = fast_pow<S>(a, K, mul, e);
    for (int i = 0; i < N; i++)
    {
        for (int j = 0; j < N; j++)
        {
            cout << ans[i][j].val() << " \n"[j == N - 1];
        }
    }

    return 0;
}

#line 1 "test/Math/pow_of_matrix/yosupo-pow_of_matrix.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/pow_of_matrix

#line 1 "Other/fast_power.hpp"
#include <bits/stdc++.h>
using namespace std;

template <class S>
S fast_pow(S x, long long n, function<S(S, S)> mul, function<S()> e)
{
    assert(0 <= n);

    S ans = e();

    while (n)
    {
        if (n & 1)
        {
            ans = mul(ans, x);
        }
        x = mul(x, x);
        n >>= 1;
    }

    return ans;
}
#line 5 "test/Math/pow_of_matrix/yosupo-pow_of_matrix.cpp"
#include <atcoder/all>

using namespace std;

int main()
{
    long long N, K;
    cin >> N >> K;

    using S = vector<vector<atcoder::modint>>;
    auto mul = [&](S a, S b) -> S
    {
        S result(N, vector<atcoder::modint>(N));
        for (int i = 0; i < N; i++)
        {
            for (int j = 0; j < N; j++)
            {
                for (int k = 0; k < N; k++)
                {
                    result[i][j] += a[i][k] * b[k][j];
                }
            }
        }

        return result;
    };

    auto e = [&]() -> S
    {
        S result = S(N, vector<atcoder::modint>(N));
        for (int i = 0; i < N; i++)
        {
            result[i][i] = 1;
        }

        return result;
    };

    vector a(N, vector<atcoder::modint>(N));
    for (int i = 0; i < N; i++)
    {
        for (int j = 0; j < N; j++)
        {
            int in_a;
            cin >> in_a;
            a[i][j] = in_a;
        }
    }

    auto ans = fast_pow<S>(a, K, mul, e);
    for (int i = 0; i < N; i++)
    {
        for (int j = 0; j < N; j++)
        {
            cout << ans[i][j].val() << " \n"[j == N - 1];
        }
    }

    return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	example_00	AC	5 ms	3 MB
g++	example_01	AC	4 ms	3 MB
g++	example_02	AC	4 ms	3 MB
g++	frobenius_hack_00	AC	4 ms	3 MB
g++	lowrank_max_random_00	AC	1438 ms	4 MB
g++	lowrank_max_random_01	AC	1317 ms	4 MB
g++	max_random_00	AC	1448 ms	4 MB
g++	max_random_01	AC	1417 ms	4 MB
g++	max_random_02	AC	1516 ms	4 MB
g++	max_random_03	AC	1433 ms	4 MB
g++	max_random_worst_00	AC	1896 ms	4 MB
g++	max_random_worst_01	AC	1909 ms	4 MB
g++	max_random_worst_02	AC	1900 ms	4 MB
g++	max_random_worst_03	AC	1897 ms	4 MB
g++	nontrivial_frobenius_form_00	AC	1378 ms	4 MB
g++	nontrivial_frobenius_form_01	AC	1456 ms	4 MB
g++	nontrivial_frobenius_form_02	AC	1498 ms	4 MB
g++	nontrivial_frobenius_form_03	AC	1425 ms	4 MB
g++	nontrivial_frobenius_form_04	AC	1307 ms	4 MB
g++	nontrivial_frobenius_form_05	AC	1387 ms	4 MB
g++	nontrivial_frobenius_form_06	AC	1443 ms	4 MB
g++	nontrivial_frobenius_form_07	AC	1341 ms	4 MB
g++	nontrivial_frobenius_form_08	AC	1514 ms	4 MB
g++	nontrivial_frobenius_form_09	AC	1485 ms	4 MB
g++	perm_max_random_00	AC	1385 ms	4 MB
g++	perm_max_random_01	AC	1216 ms	4 MB
g++	random_00	AC	1172 ms	4 MB
g++	random_01	AC	1406 ms	4 MB
g++	random_02	AC	133 ms	4 MB
g++	signed_overflow_00	AC	23 ms	4 MB
g++	small_00	AC	4 ms	3 MB
g++	small_01	AC	4 ms	3 MB
g++	small_02	AC	4 ms	3 MB
g++	small_03	AC	4 ms	3 MB
g++	small_04	AC	4 ms	3 MB
g++	small_05	AC	4 ms	3 MB
g++	small_06	AC	4 ms	3 MB
g++	small_07	AC	4 ms	3 MB
g++	small_08	AC	4 ms	3 MB
g++	small_09	AC	4 ms	3 MB
g++	small_10	AC	4 ms	3 MB
g++	small_11	AC	4 ms	3 MB
g++	unsigned_overflow_00	AC	7 ms	3 MB