test/src/linear_algebra/circulant_matrix/arc139_e.test.cpp

• View this file on GitHub
• Last update: 2024-01-30 20:57:42+09:00
• Problem: https://atcoder.jp/contests/arc139/tasks/arc139_e

Depends on

• Circulant Matrix (巡回行列) (library/linear_algebra/circulant_matrix.hpp)
• 階乗テーブル (library/math/factorial.hpp)

Code

#define PROBLEM "https://atcoder.jp/contests/arc139/tasks/arc139_e"

#include <iostream>
#include <vector>

#include <atcoder/modint>
#include <atcoder/convolution>

using mint = atcoder::modint998244353;

#include "library/math/factorial.hpp"
#include "library/linear_algebra/circulant_matrix.hpp"

using suisen::factorial;
using suisen::CirculantMatrix;

void solve(long long n, long long m) {
    if (n % 2 == 0 and m % 2 == 0) {
        std::cout << 2 << std::endl;
    } else if (n % 2 == 0) {
        factorial<mint> fac(n + 1);
        mint ans = 0;
        for (int k = 0; k <= n; ++k) if ((n - 2 * k) % m == 0) ans += fac.binom(n, k);
        std::cout << (ans * m).val() << std::endl;
    } else {
        CirculantMatrix<mint>::set_multiplication([](const std::vector<mint>& a, const std::vector<mint>& b) { return atcoder::convolution(a, b); });
        if (m % 2 == 1 and n > m) std::swap(n, m);
        std::vector<mint> dat(n);
        dat[1] = dat[n - 1] = 1;
        std::vector<mint> x(n);
        x[0] = 1;
        std::cout << ((CirculantMatrix<mint>{dat}.pow(m) * x)[0] * n).val() << std::endl;
    }
}

int main() {
    long long n, m;
    std::cin >> n >> m;
    solve(n, m);
    return 0;
}

#line 1 "test/src/linear_algebra/circulant_matrix/arc139_e.test.cpp"
#define PROBLEM "https://atcoder.jp/contests/arc139/tasks/arc139_e"

#include <iostream>
#include <vector>

#include <atcoder/modint>
#include <atcoder/convolution>

using mint = atcoder::modint998244353;

#line 1 "library/math/factorial.hpp"



#include <cassert>
#line 6 "library/math/factorial.hpp"

namespace suisen {
    template <typename T, typename U = T>
    struct factorial {
        factorial() = default;
        factorial(int n) { ensure(n); }

        static void ensure(const int n) {
            int sz = _fac.size();
            if (n + 1 <= sz) return;
            int new_size = std::max(n + 1, sz * 2);
            _fac.resize(new_size), _fac_inv.resize(new_size);
            for (int i = sz; i < new_size; ++i) _fac[i] = _fac[i - 1] * i;
            _fac_inv[new_size - 1] = U(1) / _fac[new_size - 1];
            for (int i = new_size - 1; i > sz; --i) _fac_inv[i - 1] = _fac_inv[i] * i;
        }

        T fac(const int i) {
            ensure(i);
            return _fac[i];
        }
        T operator()(int i) {
            return fac(i);
        }
        U fac_inv(const int i) {
            ensure(i);
            return _fac_inv[i];
        }
        U binom(const int n, const int r) {
            if (n < 0 or r < 0 or n < r) return 0;
            ensure(n);
            return _fac[n] * _fac_inv[r] * _fac_inv[n - r];
        }
        template <typename ...Ds, std::enable_if_t<std::conjunction_v<std::is_integral<Ds>...>, std::nullptr_t> = nullptr>
        U polynom(const int n, const Ds& ...ds) {
            if (n < 0) return 0;
            ensure(n);
            int sumd = 0;
            U res = _fac[n];
            for (int d : { ds... }) {
                if (d < 0 or d > n) return 0;
                sumd += d;
                res *= _fac_inv[d];
            }
            if (sumd > n) return 0;
            res *= _fac_inv[n - sumd];
            return res;
        }
        U perm(const int n, const int r) {
            if (n < 0 or r < 0 or n < r) return 0;
            ensure(n);
            return _fac[n] * _fac_inv[n - r];
        }
    private:
        static std::vector<T> _fac;
        static std::vector<U> _fac_inv;
    };
    template <typename T, typename U>
    std::vector<T> factorial<T, U>::_fac{ 1 };
    template <typename T, typename U>
    std::vector<U> factorial<T, U>::_fac_inv{ 1 };
} // namespace suisen


#line 1 "library/linear_algebra/circulant_matrix.hpp"



#line 7 "library/linear_algebra/circulant_matrix.hpp"

namespace suisen {
    template <typename T>
    struct CirculantMatrix {
        using value_type = T;
        using convolution_t = std::vector<value_type>(*)(const std::vector<value_type>&, const std::vector<value_type>&);

        // empty matrix
        CirculantMatrix() : CirculantMatrix(std::vector<value_type>{}) {}

        /**
         * +-                        -+
         * | a[0] a[4] a[3] a[2] a[1] |
         * | a[1] a[0] a[4] a[3] a[2] |
         * | a[2] a[1] a[0] a[4] a[3] |
         * | a[3] a[2] a[1] a[0] a[4] |
         * | a[4] a[3] a[2] a[1] a[0] |
         * +-                        -+
         */
        explicit CirculantMatrix(const std::vector<value_type>& a) : _dat(a) {}

        static void set_multiplication(convolution_t multiplication) {
            convolve = multiplication;
        }

        static CirculantMatrix<value_type> e0(int n, const value_type& zero = value_type{ 0 }) {
            return CirculantMatrix<value_type>{ std::vector<value_type>(n, zero) };
        }
        static CirculantMatrix<value_type> e1(int n, const value_type& zero = value_type{ 0 }, const value_type& one = value_type{ 1 }) {
            auto dat = std::vector<value_type>(n, zero);
            dat[0] = one;
            return CirculantMatrix<value_type>{ dat };
        }

        int size() const {
            return _dat.size();
        }

        value_type get(int i, int j) const {
            const int n = size();
            int k = i - j;
            if (k < 0) k += n;
            return _dat[k];
        }
        value_type operator[](const std::pair<int, int> &p) const {
            return get(p.first, p.second);
        }

        friend CirculantMatrix<value_type> operator+(const CirculantMatrix<value_type>& mat) {
            return mat;
        }
        friend CirculantMatrix<value_type> operator-(const CirculantMatrix<value_type>& mat) {
            const int n = mat.size();
            std::vector<value_type> res(n);
            for (int i = 0; i < n; ++i) res[i] = -mat._dat[i];
            return CirculantMatrix<value_type> { std::move(res) };
        }
        friend CirculantMatrix<value_type> operator+(const CirculantMatrix<value_type>& lhs, const CirculantMatrix<value_type>& rhs) {
            const int n = lhs.size();
            assert(n == int(rhs.size()));
            std::vector<value_type> res(n);
            for (int i = 0; i < n; ++i) res[i] = lhs._dat[i] + rhs._dat[i];
            return CirculantMatrix<value_type> { std::move(res) };
        }
        friend CirculantMatrix<value_type> operator-(const CirculantMatrix<value_type>& lhs, const CirculantMatrix<value_type>& rhs) {
            const int n = lhs.size();
            assert(n == int(rhs.size()));
            std::vector<value_type> res(n);
            for (int i = 0; i < n; ++i) res[i] = lhs._dat[i] - rhs._dat[i];
            return CirculantMatrix<value_type> { std::move(res) };
        }
        friend CirculantMatrix<value_type> operator*(const CirculantMatrix<value_type>& lhs, const CirculantMatrix<value_type>& rhs) {
            const int n = lhs.size();
            assert(n == int(rhs.size()));
            std::vector<value_type> res = convolve(lhs._dat, rhs._dat);
            for (int i = n; i < int(res.size()); ++i) res[i - n] += res[i];
            res.resize(n);
            return CirculantMatrix<value_type> { std::move(res) };
        }
        friend std::vector<value_type> operator*(const CirculantMatrix<value_type>& mat, const std::vector<value_type>& x) {
            return std::move((mat * CirculantMatrix<value_type> { x })._dat);
        }
        friend CirculantMatrix<value_type> operator*(const CirculantMatrix<value_type>& mat, const value_type& coef) {
            const int n = mat.size();
            std::vector<value_type> res(n);
            for (int i = 0; i < n; ++i) res[i] = coef * mat._dat[i];
            return CirculantMatrix<value_type> { res };
        }
        friend CirculantMatrix<value_type> operator*(const value_type& coef, const CirculantMatrix<value_type>& mat) {
            return mat * coef;
        }

        CirculantMatrix<value_type>& operator+=(const CirculantMatrix<value_type>& rhs) {
            return *this = *this + rhs;
        }
        CirculantMatrix<value_type>& operator-=(const CirculantMatrix<value_type>& rhs) {
            return *this = *this - rhs;
        }
        CirculantMatrix<value_type>& operator*=(const CirculantMatrix<value_type>& rhs) {
            return *this = *this * rhs;
        }
        CirculantMatrix<value_type>& operator*=(const value_type& coef) {
            return *this = *this * coef;
        }

        CirculantMatrix<value_type> pow(long long b) {
            auto res = CirculantMatrix<value_type>::e1(size());
            for (auto p = *this; b; b >>= 1) {
                if (b & 1) res *= p;
                p *= p;
            }
            return res;
        }

    private:
        static inline convolution_t convolve{
            [](const auto&, const auto&) {
                std::cerr << "convolution function is not available." << std::endl;
                assert(false);
                return std::vector<value_type>{};
            }
        };

        std::vector<value_type> _dat;
    };
} // namespace suisen


#line 13 "test/src/linear_algebra/circulant_matrix/arc139_e.test.cpp"

using suisen::factorial;
using suisen::CirculantMatrix;

void solve(long long n, long long m) {
    if (n % 2 == 0 and m % 2 == 0) {
        std::cout << 2 << std::endl;
    } else if (n % 2 == 0) {
        factorial<mint> fac(n + 1);
        mint ans = 0;
        for (int k = 0; k <= n; ++k) if ((n - 2 * k) % m == 0) ans += fac.binom(n, k);
        std::cout << (ans * m).val() << std::endl;
    } else {
        CirculantMatrix<mint>::set_multiplication([](const std::vector<mint>& a, const std::vector<mint>& b) { return atcoder::convolution(a, b); });
        if (m % 2 == 1 and n > m) std::swap(n, m);
        std::vector<mint> dat(n);
        dat[1] = dat[n - 1] = 1;
        std::vector<mint> x(n);
        x[0] = 1;
        std::cout << ((CirculantMatrix<mint>{dat}.pow(m) * x)[0] * n).val() << std::endl;
    }
}

int main() {
    long long n, m;
    std::cin >> n >> m;
    solve(n, m);
    return 0;
}

Back to top page