test/src/polynomial/compose/composition_of_formal_power_series.test.cpp

View this file on GitHub
Last update: 2023-05-27 18:49:30+09:00
Problem: https://judge.yosupo.jp/problem/composition_of_formal_power_series

Depends on

合成 (library/polynomial/compose.hpp)

Code

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

#include <iostream>

#include "library/polynomial/compose.hpp"

#include <atcoder/modint>

using mint = atcoder::modint998244353;

namespace atcoder {
    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;
    }
} // namespace atcoder

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

    int n;
    std::cin >> n;

    std::vector<mint> a(n), b(n);
    for (auto &e : a) std::cin >> e;
    for (auto &e : b) std::cin >> e;

    std::vector<mint> c = suisen::compose(a, b, n);
    for (int i = 0; i < n; ++i) {
        std::cout << c[i];
        if (i + 1 != n) std::cout << ' ';
    }
    std::cout << '\n';
}

#line 1 "test/src/polynomial/compose/composition_of_formal_power_series.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/composition_of_formal_power_series"

#include <iostream>

#line 1 "library/polynomial/compose.hpp"



#include <cmath>
#include <vector>
#include <atcoder/convolution>

namespace suisen {
    template <typename mint>
    std::vector<mint> compose(const std::vector<mint>& f, std::vector<mint> g, const int n) {
        std::vector<mint> res(n);
        if (n == 0) return res;
        if (f.empty()) return res;

        if (std::find_if(g.begin(), g.end(), [](mint x) { return x != 0; }) == g.end()) return res[0] = f[0], res;

        // taylor shift f(x + [x^0]g)
        const std::vector<mint> fa = [&]{
            const mint a = std::exchange(g[0], 0);
            const int siz_f = f.size();
            
            std::vector<mint> fac(siz_f), fac_inv(siz_f);
            fac[0] = 1;
            for (int i = 1; i <= siz_f - 1; ++i) fac[i] = fac[i - 1] * i;
            fac_inv[siz_f - 1] = fac[siz_f - 1].inv();
            for (int i = siz_f - 1; i >= 1; --i) fac_inv[i - 1] = fac_inv[i] * i;

            std::vector<mint> ec(siz_f), fa(siz_f);
            mint p = 1;
            for (int i = 0; i < siz_f; ++i, p *= a) {
                ec[i] = p * fac_inv[i];
                fa[siz_f - 1 - i] = (i < int(f.size()) ? f[i] : 0) * fac[i];
            }
            fa = atcoder::convolution(fa, ec), fa.resize(siz_f);
            std::reverse(fa.begin(), fa.end());
            for (int i = 0; i < siz_f; ++i) {
                fa[i] *= fac_inv[i];
            }
            if (siz_f > n) fa.resize(n);
            return fa;
        }();

        const int sqn = ::sqrt(f.size()) + 1;

        const int z = [n]{
            int z = 1;
            while (z < 2 * n - 1) z <<= 1;
            return z;
        }();
        const mint iz = mint(z).inv();

        g.erase(g.begin());
        g.resize(z);
        atcoder::internal::butterfly(g);

        auto mult_g = [&](std::vector<mint> a) {
            a.resize(z);
            atcoder::internal::butterfly(a);
            for (int j = 0; j < z; ++j) a[j] *= g[j] * iz;
            atcoder::internal::butterfly_inv(a);
            a.resize(n);
            return a;
        };

        std::vector<std::vector<mint>> pow_g(sqn, std::vector<mint>(n));
        pow_g[0][0] = 1;
        for (int i = 1; i < sqn; ++i) {
            pow_g[i] = mult_g(pow_g[i - 1]);
        }

        std::vector<mint> gl = mult_g(pow_g[sqn - 1]);
        gl.resize(z);
        atcoder::internal::butterfly(gl);

        std::vector<mint> pow_gl(z);
        pow_gl[0] = 1;

        for (int i = 0; i < sqn; ++i) {
            const int off_i = i * sqn;
            const int siz_i = n - off_i;
            if (siz_i < 0) break;
            std::vector<mint> fg(siz_i);
            for (int j = 0; j < sqn; ++j) {
                const int ij = i * sqn + j;
                if (ij >= int(fa.size())) break;

                const mint c = fa[ij];
                const std::vector<mint>& gj = pow_g[j];
                for (int k = 0; k < siz_i - j; ++k) {
                    fg[j + k] += c * gj[k];
                }
            }
            fg.resize(z);
            atcoder::internal::butterfly(pow_gl);
            atcoder::internal::butterfly(fg);
            for (int k = 0; k < z; ++k) {
                fg[k] *= pow_gl[k] * iz;
                pow_gl[k] *= gl[k] * iz;
            }
            atcoder::internal::butterfly_inv(pow_gl);
            atcoder::internal::butterfly_inv(fg);
            for (int k = 0; k < siz_i; ++k) {
                res[off_i + k] += fg[k];
            }
            std::fill(pow_gl.begin() + n, pow_gl.end(), 0);
        }
        return res;
    }
} // namespace suisen



#line 6 "test/src/polynomial/compose/composition_of_formal_power_series.test.cpp"

#include <atcoder/modint>

using mint = atcoder::modint998244353;

namespace atcoder {
    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;
    }
} // namespace atcoder

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

    int n;
    std::cin >> n;

    std::vector<mint> a(n), b(n);
    for (auto &e : a) std::cin >> e;
    for (auto &e : b) std::cin >> e;

    std::vector<mint> c = suisen::compose(a, b, n);
    for (int i = 0; i < n; ++i) {
        std::cout << c[i];
        if (i + 1 != n) std::cout << ' ';
    }
    std::cout << '\n';
}