cp-library-cpp

This documentation is automatically generated by online-judge-tools/verification-helper

View the Project on GitHub suisen-cp/cp-library-cpp

:heavy_check_mark: test/src/convolution/multi_variate_convolution/multivariate_convolution.test.cpp

Depends on

Code

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

#include <iostream>

#include <atcoder/modint>
using mint = atcoder::modint998244353;

std::istream& operator>>(std::istream& in, mint &a) {
    long long e; in >> e; a = e;
    return in;
}

std::ostream& operator<<(std::ostream& out, const mint &a) {
    out << a.val();
    return out;
}

#include "library/convolution/multi_variate_convolution.hpp"
using namespace suisen;

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int k;
    std::cin >> k;
    std::vector<int> dim(k);
    for (int &e : dim) std::cin >> e;

    multi_variate_convolution convolution(dim);
    int n = convolution.size();

    std::vector<mint> f(n), g(n);
    for (auto &e : f) std::cin >> e;
    for (auto &e : g) std::cin >> e;
    auto h = convolution(f, g);

    for (int i = 0; i < n; ++i) std::cout << h[i] << " \n"[i == n - 1];

    return 0;
}
#line 1 "test/src/convolution/multi_variate_convolution/multivariate_convolution.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/multivariate_convolution"

#include <iostream>

#include <atcoder/modint>
using mint = atcoder::modint998244353;

std::istream& operator>>(std::istream& in, mint &a) {
    long long e; in >> e; a = e;
    return in;
}

std::ostream& operator<<(std::ostream& out, const mint &a) {
    out << a.val();
    return out;
}

#line 1 "library/convolution/multi_variate_convolution.hpp"



#include <numeric>
#include <atcoder/convolution>

#line 1 "library/util/default_operator.hpp"



namespace suisen {
    namespace default_operator {
        template <typename T>
        auto zero() -> decltype(T { 0 }) { return T { 0 }; }
        template <typename T>
        auto one()  -> decltype(T { 1 }) { return T { 1 }; }
        template <typename T>
        auto add(const T &x, const T &y) -> decltype(x + y) { return x + y; }
        template <typename T>
        auto sub(const T &x, const T &y) -> decltype(x - y) { return x - y; }
        template <typename T>
        auto mul(const T &x, const T &y) -> decltype(x * y) { return x * y; }
        template <typename T>
        auto div(const T &x, const T &y) -> decltype(x / y) { return x / y; }
        template <typename T>
        auto mod(const T &x, const T &y) -> decltype(x % y) { return x % y; }
        template <typename T>
        auto neg(const T &x) -> decltype(-x) { return -x; }
        template <typename T>
        auto inv(const T &x) -> decltype(one<T>() / x)  { return one<T>() / x; }
    } // default_operator
    namespace default_operator_noref {
        template <typename T>
        auto zero() -> decltype(T { 0 }) { return T { 0 }; }
        template <typename T>
        auto one()  -> decltype(T { 1 }) { return T { 1 }; }
        template <typename T>
        auto add(T x, T y) -> decltype(x + y) { return x + y; }
        template <typename T>
        auto sub(T x, T y) -> decltype(x - y) { return x - y; }
        template <typename T>
        auto mul(T x, T y) -> decltype(x * y) { return x * y; }
        template <typename T>
        auto div(T x, T y) -> decltype(x / y) { return x / y; }
        template <typename T>
        auto mod(T x, T y) -> decltype(x % y) { return x % y; }
        template <typename T>
        auto neg(T x) -> decltype(-x) { return -x; }
        template <typename T>
        auto inv(T x) -> decltype(one<T>() / x)  { return one<T>() / x; }
    } // default_operator
} // namespace suisen


#line 8 "library/convolution/multi_variate_convolution.hpp"

namespace suisen {
    struct multi_variate_convolution {
        multi_variate_convolution() : multi_variate_convolution(std::vector<int>{}) {}
        multi_variate_convolution(const std::vector<int> &dim) : _n(std::accumulate(dim.begin(), dim.end(), 1, std::multiplies<int>())), _k(dim.size()), _m(2 * ceil_pow2(_n)), _chi(_n, 0) {
            for (int i = 0; i < _n; ++i) {
                int den = 1;
                for (int e : dim) den *= e, _chi[i] += i / den;
                if (_k) _chi[i] %= _k;
            }
        }

        int size() const { return _n; }
        int dim_num() const { return _k; }

        template <typename T, typename Inverse = decltype(default_operator::inv<T>)>
        std::vector<T> convolution(std::vector<T> f, std::vector<T> g, Inverse inv = default_operator::inv<T>) const {
            assert(int(f.size()) == _n and int(g.size()) == _n);
            if (_k == 0) return { f[0] * g[0] };
            auto rf = ranked(f), rg = ranked(g);
            for (auto &v : rf) atcoder::internal::butterfly(v);
            for (auto &v : rg) atcoder::internal::butterfly(v);
            std::vector rh(_k, std::vector<T>(_m, T{}));
            for (int i = 0; i < _k; ++i) for (int j = 0; j < _k; ++j) {
                int r = i + j < _k ? i + j : i + j - _k;
                for (int p = 0; p < _m; ++p) rh[r][p] += rf[i][p] * rg[j][p];
            }
            for (auto &v : rh) atcoder::internal::butterfly_inv(v);
            const T isz = inv(T(_m));
            std::vector<T> h(_n);
            for (int i = 0; i < _n; ++i) h[i] = rh[_chi[i]][i] * isz;
            return h;
        }

        template <typename T>
        std::vector<T> operator()(std::vector<T> f, std::vector<T> g) const {
            return convolution(f, g);
        }

    private:
        int _n, _k, _m;
        std::vector<int> _chi;

        static constexpr int ceil_pow2(int m) {
            int res = 1;
            while (res < m) res *= 2;
            return res;
        }

        template <typename T>
        auto ranked(const std::vector<T> &f) const {
            std::vector rf(_k, std::vector<T>(_m, T{}));
            for (int i = 0; i < _n; ++i) rf[_chi[i]][i] = f[i];
            return rf;
        }
    };
} // namespace suisen



#line 19 "test/src/convolution/multi_variate_convolution/multivariate_convolution.test.cpp"
using namespace suisen;

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);

    int k;
    std::cin >> k;
    std::vector<int> dim(k);
    for (int &e : dim) std::cin >> e;

    multi_variate_convolution convolution(dim);
    int n = convolution.size();

    std::vector<mint> f(n), g(n);
    for (auto &e : f) std::cin >> e;
    for (auto &e : g) std::cin >> e;
    auto h = convolution(f, g);

    for (int i = 0; i < n; ++i) std::cout << h[i] << " \n"[i == n - 1];

    return 0;
}
Back to top page