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/polynomial/compose/dummy.test.cpp

Depends on

Code

#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"

#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

std::vector<mint> naive(std::vector<mint> f, std::vector<mint> g, int n) {
    if (n == 0) return {};

    std::vector<mint> powg(n);
    powg[0] = 1;

    std::vector<mint> fg(n);
    for (mint fi : f) {
        for (int j = 0; j < n; ++j) {
            fg[j] += fi * powg[j];
        }
        powg = atcoder::convolution(powg, g);
        powg.resize(n);
    }
    return fg;
}

void test() {
    {
        std::vector<mint> f { 1, 2, 3, 4 };
        std::vector<mint> g { 5, 4, 3, 2 };
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f { 1, 2, 3, 4, 5, 6, 7 };
        std::vector<mint> g { 5, 4, 3, 2 };
        assert((suisen::compose(f, g, 8) == naive(f, g, 8)));
        assert((suisen::compose(f, g, 7) == naive(f, g, 7)));
        assert((suisen::compose(f, g, 6) == naive(f, g, 6)));
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f {};
        std::vector<mint> g { 5, 4, 3 };
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f { 2 };
        std::vector<mint> g { 5, 4, 3 };
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f { 2 };
        std::vector<mint> g { 3 };
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f { 5, 4, 3 };
        std::vector<mint> g { 2 };
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f { 5, 4, 3 };
        std::vector<mint> g {};
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f {};
        std::vector<mint> g {};
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
}

int main() {
    test();
    std::cout << "Hello World" << std::endl;
    return 0;
}
#line 1 "test/src/polynomial/compose/dummy.test.cpp"
#define PROBLEM "https://judge.u-aizu.ac.jp/onlinejudge/description.jsp?id=ITP1_1_A"

#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/dummy.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

std::vector<mint> naive(std::vector<mint> f, std::vector<mint> g, int n) {
    if (n == 0) return {};

    std::vector<mint> powg(n);
    powg[0] = 1;

    std::vector<mint> fg(n);
    for (mint fi : f) {
        for (int j = 0; j < n; ++j) {
            fg[j] += fi * powg[j];
        }
        powg = atcoder::convolution(powg, g);
        powg.resize(n);
    }
    return fg;
}

void test() {
    {
        std::vector<mint> f { 1, 2, 3, 4 };
        std::vector<mint> g { 5, 4, 3, 2 };
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f { 1, 2, 3, 4, 5, 6, 7 };
        std::vector<mint> g { 5, 4, 3, 2 };
        assert((suisen::compose(f, g, 8) == naive(f, g, 8)));
        assert((suisen::compose(f, g, 7) == naive(f, g, 7)));
        assert((suisen::compose(f, g, 6) == naive(f, g, 6)));
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f {};
        std::vector<mint> g { 5, 4, 3 };
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f { 2 };
        std::vector<mint> g { 5, 4, 3 };
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f { 2 };
        std::vector<mint> g { 3 };
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f { 5, 4, 3 };
        std::vector<mint> g { 2 };
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f { 5, 4, 3 };
        std::vector<mint> g {};
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
    {
        std::vector<mint> f {};
        std::vector<mint> g {};
        assert((suisen::compose(f, g, 5) == naive(f, g, 5)));
        assert((suisen::compose(f, g, 4) == naive(f, g, 4)));
        assert((suisen::compose(f, g, 3) == naive(f, g, 3)));
        assert((suisen::compose(f, g, 2) == naive(f, g, 2)));
        assert((suisen::compose(f, g, 1) == naive(f, g, 1)));
        assert((suisen::compose(f, g, 0) == naive(f, g, 0)));
    }
}

int main() {
    test();
    std::cout << "Hello World" << std::endl;
    return 0;
}
Back to top page