Formal Power Series Relaxed
(library/polynomial/formal_power_series_relaxed.hpp)

View this file on GitHub
Last update: 2023-07-09 04:04:16+09:00
Include: #include "library/polynomial/formal_power_series_relaxed.hpp"

Formal Power Series Relaxed

形式的べき級数 $f\in \mathbb{K}\lbrack\lbrack x\rbrack\rbrack$ の $x ^ i$ の係数を $f_i$ と書く．

以下に示す漸化式に従って Relaxed Convolution を適用することで，$\Theta(N (\log N) ^ 2)$ 時間で $\mathrm{inv},\log,\exp$ などをオンラインで計算できる．

$h = f / g$ (未実装)

$g _ 0 \in \mathbb{K} ^ {\times}$ を仮定する．

\[\begin{aligned} f _ n &= \sum _ {i = 0} ^ n h _ i g _ {n - i}\\ &= h _ n g _ 0 + \sum _ {i = 0} ^ {n - 1} h _ i g _ {n - i},\\ h _ n &= \dfrac{1}{g _ 0}\left(f _ n - \sum _ {i = 0} ^ {n - 1}h _ i g _ {n - i} \right). \end{aligned}\]

$g = 1/f$

$f _ 0 \in \mathbb{K} ^ {\times}$ を仮定する．

\[g _ n = \begin{cases} f _ 0 ^ {-1} & \text{if } n = 0\\ \displaystyle -f _ 0 ^ {-1} \sum _ {i = 0} ^ {n - 1} g _ i f _ {n - i} & \text{otherwise} \end{cases}\]

$g = \exp f$

$f _ 0 = 0$ を仮定する．$g’ = f’ g$ より，

\[\begin{aligned} (n + 1) g _ {n + 1} &= \sum _ {i = 0} ^ n (n - i + 1) f _ {n - i + 1} g _ i,\\ g _ n &= \begin{cases} 1 & \text{if }n = 0 \\ \dfrac{1}{n}\displaystyle\sum _ {i = 0} ^ {n - 1} (n - i) f _ {n - i } g _ i &\text{otherwise} \end{cases} \end{aligned}\]

$g = \log f$

$g _ 0 = 1$ を仮定する．$fg’=f’$ より，

\[\begin{aligned} \sum _ {i = 0} ^ n (i+1) g _ {i + 1} f _ {n - i} &= (n + 1)f _ {n + 1},\\ (n + 1) f _ 0 g _ {n + 1} &= (n + 1) f _ {n + 1} - \sum _ {i = 0} ^ {n - 1}(i+1) g _ {i + 1} f _ {n - i},\\ g _ {n + 1} &= \dfrac{1}{(n+1)f_0}\left((n + 1) f _ {n + 1} - \sum _ {i = 0} ^ {n - 1}(i+1) g _ {i + 1} f _ {n - i}\right),\\ g _ n &= \begin{cases} 0 & \text{if }n = 0\\ \displaystyle \dfrac{1}{nf_0}\left(n f _ {n} - \sum _ {i = 1} ^ {n - 1}i g _ i f _ {n - i}\right) & \text{otherwise} \end{cases} \end{aligned}\]

$g = f ^ k$

$f_0\neq 0$ を仮定する．$f _ 0 = 0$ の場合は，$f = x ^ p \cdot f’\ (f’_0 \neq 0)$ の形に直してから $g = x ^ {pk} f’^ k$ として計算すればよい．$g’=kf^{k-1} f’$ の両辺に $f$ を掛けることで $fg’=kgf’$ が得られるので，

\[\begin{aligned} \sum _ {i = 0} ^ n (i + 1) g _ {i + 1} f _ {n - i} &= k\sum _ {i = 0} ^ n g _ i\cdot (n-i+1) f _ {n-i+1}, \\ (n+1)f_0g _ {n+1} &= k\sum _ {i = 0} ^ n g _ i\cdot (n-i+1) f _ {n-i+1} - \sum _ {i = 1} ^ {n} i g _ i f _ {n - i + 1} \\ &=\sum _ {i = 0} ^ {n} (k(n-i+1)-i)g_i f _{n-i+1},\\ g _ {n + 1} &= \dfrac{1}{(n+1)f_0}\sum _ {i = 0} ^ {n} (k(n-i+1)-i)g_i f _{n-i+1},\\ g _ n &= \begin{cases} f _ 0 ^ k & \text{if }n = 0\\ \displaystyle\dfrac{1}{nf_0}\sum _ {i = 0} ^ {n-1} (k(n-i)-i)g_i f _{n-i} & \text{otherwise} \end{cases} \end{aligned}\]

$g = f ^ {1/2}$

$f _ 0 \neq 0$ であることおよび $f _ 0 ^ {1/2}$ の存在を仮定する．

\[\begin{aligned} f _ {n+1} &= \sum _ {i = 0} ^ {n + 1} g _ i g _ {n + 1 - i},\\ 2g _ 0 g _ {n + 1} &= f _ {n + 1} - \sum _ {i = 1} ^ {n} g _ i g _ {n + 1 - i},\\ g _ n &= \begin{cases} f _ 0 ^ {1/2} & \text{if }n = 0\\ \displaystyle \dfrac{1}{2g_0}\left(f _ {n } - \sum _ {i = 1} ^ {n-1} g _ i g _ {n - i}\right) & \text{otherwise} \end{cases} \end{aligned}\]

$g = f ^ {1/k}$ (未実装)

$f _ 0 \neq 0$ であることおよび $f _ 0 ^ {1/k}$ の存在を仮定する．$g ^ k = f$ を $\mathrm{pow}$ と同様に微分することで，次を得る．

\[g _ n = \begin{cases} f _ 0 ^ {1/k} & \text{if }n = 0\\ \displaystyle\dfrac{1}{nf_0}\sum _ {i = 0} ^ {n-1} (k^{-1}(n-i)-i)g_i f _{n-i} & \text{otherwise} \end{cases}\]

より一般には，$g = f ^ {p/q}$ を満たす $g$ を計算できる．

Depends on

Verified with

Code

#ifndef SUISEN_FPS_RELAXED
#define SUISEN_FPS_RELAXED

#include <atcoder/convolution>
#include "library/math/inv_mods.hpp"
#include "library/convolution/relaxed_convolution_ntt.hpp"
#include "library/math/modint_extension.hpp"

namespace suisen {
    template <typename mint>
    struct RelaxedInv {
        mint append(const mint& fi) {
            const int i = g.size();
            if (i == 0) {
                assert(fi != 0);
                g.push_back(fi.inv());
            } else {
                g.push_back(-g[0] * fg.append(fi, g[i - 1]));
            }
            return g.back();
        }
        mint operator[](int i) const {
            return g[i];
        }
    private:
        std::vector<mint> g;
        RelaxedConvolutionNTT<mint> fg;
    };

    template <typename mint>
    struct RelaxedExp {
        mint append(const mint& fi) {
            static inv_mods<mint> invs;
            const int i = g.size();
            if (i == 0) {
                assert(fi == 0);
                g.push_back(1);
            } else {
                g.push_back(df_g.append(i * fi, g[i - 1]) * invs[i]);
            }
            return g.back();
        }
        mint operator[](int i) const {
            return g[i];
        }
    private:
        std::vector<mint> g;
        RelaxedConvolutionNTT<mint> df_g;
    };

    template <typename mint>
    struct RelaxedLog {
        mint append(const mint& fi) {
            static inv_mods<mint> invs;
            f.push_back(fi);
            const int i = g.size();
            if (i == 0) {
                assert(f[i] == 1);
                g.push_back(0);
            } else if (i == 1) {
                g.push_back(f[i]);
            } else {
                g.push_back(f[i] - fg.append((i - 1) * g[i - 1], f[i - 1]) * invs[i]);
            }
            return g.back();
        }
        mint operator[](int i) const {
            return g[i];
        }
    private:
        std::vector<mint> f, g;
        RelaxedConvolutionNTT<mint> fg;
    };

    template <typename mint>
    struct RelaxedPow {
        RelaxedPow(long long k = 0) : k(k) {}

        mint append(const mint& fi) {
            if (k == 0) {
                return g.emplace_back(g.empty() ? 1 : 0);
            }
            static inv_mods<mint> invs;
            if (is_zero) {
                if (fi == 0) {
                    z = std::min(z + k, 1000000000LL);
                } else {
                    is_zero = false;
                    inv_base = fi.inv();
                }
            }
            if (not is_zero) {
                f.push_back(fi);
            }
            if (index < z) {
                g.push_back(0);
            } else if (index == z) {
                g.push_back(f[0].pow(k));
            } else {
                int i = index - z;
                mint v1 = fg1.append(mint(k - (i - 1)) * g[z + i - 1], f[i]);
                mint v2 = fg2.append(g[z + i - 1], mint(k) * (i - 1) * f[i]);
                g.push_back((v1 + v2) * inv_base * invs[i]);
            }
            ++index;
            return g.back();
        }
        mint operator[](int i) const {
            return g[i];
        }
    private:
        long long k;
        long long z = 0;
        long long index = 0;
        bool is_zero = true;
        mint inv_base = 0;

        std::vector<mint> f, g;
        RelaxedConvolutionNTT<mint> fg1;
        RelaxedConvolutionNTT<mint> fg2;
    };

    template <typename mint>
    struct RelaxedSqrt {
        std::optional<mint> append(const mint& fi) {
            if (g.empty()) {
                auto opt_g0 = safe_sqrt(fi);
                if (not opt_g0) return std::nullopt;
                mint g0 = *opt_g0;
                c = (2 * g0).inv();
                return g.emplace_back(g0);
            } else if (g.size() == 1) {
                return g.emplace_back(c * fi);
            } else {
                mint gi = c * (fi - gg.append(g.back(), g.back()));
                return g.emplace_back(gi);
            }
        }
        mint operator[](int i) const {
            return g[i];
        }
    private:
        mint c = 0;
        std::vector<mint> g;
        RelaxedConvolutionNTT<mint> gg;
    };
} // namespace suisen


#endif // SUISEN_FPS_RELAXED

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



#include <atcoder/convolution>
#line 1 "library/math/inv_mods.hpp"



#include <vector>

namespace suisen {
    template <typename mint>
    class inv_mods {
    public:
        inv_mods() = default;
        inv_mods(int n) { ensure(n); }
        const mint& operator[](int i) const {
            ensure(i);
            return invs[i];
        }
        static void ensure(int n) {
            int sz = invs.size();
            if (sz < 2) invs = { 0, 1 }, sz = 2;
            if (sz < n + 1) {
                invs.resize(n + 1);
                for (int i = sz; i <= n; ++i) invs[i] = mint(mod - mod / i) * invs[mod % i];
            }
        }
    private:
        static std::vector<mint> invs;
        static constexpr int mod = mint::mod();
    };
    template <typename mint>
    std::vector<mint> inv_mods<mint>::invs{};

    template <typename mint>
    std::vector<mint> get_invs(const std::vector<mint>& vs) {
        const int n = vs.size();

        mint p = 1;
        for (auto& e : vs) {
            p *= e;
            assert(e != 0);
        }
        mint ip = p.inv();

        std::vector<mint> rp(n + 1);
        rp[n] = 1;
        for (int i = n - 1; i >= 0; --i) {
            rp[i] = rp[i + 1] * vs[i];
        }
        std::vector<mint> res(n);
        for (int i = 0; i < n; ++i) {
            res[i] = ip * rp[i + 1];
            ip *= vs[i];
        }
        return res;
    }
}


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



#line 5 "library/convolution/relaxed_convolution_ntt.hpp"

namespace suisen {
    // reference: https://qiita.com/Kiri8128/items/1738d5403764a0e26b4c
    template <typename mint>
    struct RelaxedConvolutionNTT {
        RelaxedConvolutionNTT(): _n(0), _f{}, _g{}, _h{} {}

        mint append(const mint& fi, const mint& gi) {
            static constexpr int threshold_log = 6;
            static constexpr int threshold = 1 << threshold_log;
            static constexpr int threshold_mask = threshold - 1;

            ++_n;
            _f.push_back(fi), _g.push_back(gi);

            const int q = _n >> threshold_log, r = _n & threshold_mask;
            if (r == 0) {
                if (q == (-q & q)) {
                    std::vector<mint> f_fft = _f;
                    std::vector<mint> g_fft = _g;
                    f_fft.resize(2 * _n);
                    g_fft.resize(2 * _n);
                    atcoder::internal::butterfly(f_fft);
                    atcoder::internal::butterfly(g_fft);
                    std::vector<mint> h(2 * _n);
                    for (int i = 0; i < 2 * _n; ++i) {
                        h[i] = f_fft[i] * g_fft[i];
                    }
                    atcoder::internal::butterfly_inv(h);
                    ensure(2 * _n);
                    const mint z = mint(2 * _n).inv();
                    for (int i = _n - 1; i < 2 * _n; ++i) {
                        _h[i] += h[i] * z;
                    }
                    _f_fft.push_back(std::move(f_fft));
                    _g_fft.push_back(std::move(g_fft));
                } else {
                    const int log_q = __builtin_ctz(q);
                    const int k = (-q & q) << threshold_log;

                    std::vector<mint> f_fft(_f.end() - k, _f.end());
                    std::vector<mint> g_fft(_g.end() - k, _g.end());
                    f_fft.resize(2 * k);
                    g_fft.resize(2 * k);
                    atcoder::internal::butterfly(f_fft);
                    atcoder::internal::butterfly(g_fft);
                    std::vector<mint> h(2 * k);
                    for (int i = 0; i < 2 * k; ++i) {
                        h[i] = _f_fft[log_q + 1][i] * g_fft[i] + f_fft[i] * _g_fft[log_q + 1][i];
                    }
                    atcoder::internal::butterfly_inv(h);
                    const mint z = mint(2 * k).inv();
                    for (int i = 0; i < k; ++i) {
                        _h[_n - 1 + i] += h[k - 1 + i] * z;
                    }
                }
            } else {
                // naive convolve
                ensure(_n);
                for (int i = 0; i < r; ++i) {
                    _h[_n - 1] += _f[i] * _g[_n - 1 - i];
                }
                if (_n != r) {
                    for (int i = 0; i < r; ++i) {
                        _h[_n - 1] += _f[_n - i - 1] * _g[i];
                    }
                }
            }
            return _h[_n - 1];
        }

        const mint& operator[](int i) const {
            return _h[i];
        }
        std::vector<mint> get() const {
            return _h;
        }

    private:
        int _n;
        std::vector<mint> _f, _g, _h;

        std::vector<std::vector<mint>> _f_fft, _g_fft;

        void ensure(std::size_t n) {
            if (_h.size() < n) _h.resize(n);
        }
    };
} // namespace suisen



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



#include <cassert>
#include <optional>

/**
 * refernce: https://37zigen.com/tonelli-shanks-algorithm/
 * calculates x s.t. x^2 = a mod p in O((log p)^2).
 */
template <typename mint>
std::optional<mint> safe_sqrt(mint a) {
    static int p = mint::mod();
    if (a == 0) return std::make_optional(0);
    if (p == 2) return std::make_optional(a);
    if (a.pow((p - 1) / 2) != 1) return std::nullopt;
    mint b = 1;
    while (b.pow((p - 1) / 2) == 1) ++b;
    static int tlz = __builtin_ctz(p - 1), q = (p - 1) >> tlz;
    mint x = a.pow((q + 1) / 2);
    b = b.pow(q);
    for (int shift = 2; x * x != a; ++shift) {
        mint e = a.inv() * x * x;
        if (e.pow(1 << (tlz - shift)) != 1) x *= b;
        b *= b;
    }
    return std::make_optional(x);
}

/**
 * calculates x s.t. x^2 = a mod p in O((log p)^2).
 * if not exists, raises runtime error.
 */
template <typename mint>
auto sqrt(mint a) -> decltype(mint::mod(), mint()) {
    return *safe_sqrt(a);
}
template <typename mint>
auto log(mint a) -> decltype(mint::mod(), mint()) {
    assert(a == 1);
    return 0;
}
template <typename mint>
auto exp(mint a) -> decltype(mint::mod(), mint()) {
    assert(a == 0);
    return 1;
}
template <typename mint, typename T>
auto pow(mint a, T b) -> decltype(mint::mod(), mint()) {
    return a.pow(b);
}
template <typename mint>
auto inv(mint a) -> decltype(mint::mod(), mint()) {
    return a.inv();
}


#line 8 "library/polynomial/formal_power_series_relaxed.hpp"

namespace suisen {
    template <typename mint>
    struct RelaxedInv {
        mint append(const mint& fi) {
            const int i = g.size();
            if (i == 0) {
                assert(fi != 0);
                g.push_back(fi.inv());
            } else {
                g.push_back(-g[0] * fg.append(fi, g[i - 1]));
            }
            return g.back();
        }
        mint operator[](int i) const {
            return g[i];
        }
    private:
        std::vector<mint> g;
        RelaxedConvolutionNTT<mint> fg;
    };

    template <typename mint>
    struct RelaxedExp {
        mint append(const mint& fi) {
            static inv_mods<mint> invs;
            const int i = g.size();
            if (i == 0) {
                assert(fi == 0);
                g.push_back(1);
            } else {
                g.push_back(df_g.append(i * fi, g[i - 1]) * invs[i]);
            }
            return g.back();
        }
        mint operator[](int i) const {
            return g[i];
        }
    private:
        std::vector<mint> g;
        RelaxedConvolutionNTT<mint> df_g;
    };

    template <typename mint>
    struct RelaxedLog {
        mint append(const mint& fi) {
            static inv_mods<mint> invs;
            f.push_back(fi);
            const int i = g.size();
            if (i == 0) {
                assert(f[i] == 1);
                g.push_back(0);
            } else if (i == 1) {
                g.push_back(f[i]);
            } else {
                g.push_back(f[i] - fg.append((i - 1) * g[i - 1], f[i - 1]) * invs[i]);
            }
            return g.back();
        }
        mint operator[](int i) const {
            return g[i];
        }
    private:
        std::vector<mint> f, g;
        RelaxedConvolutionNTT<mint> fg;
    };

    template <typename mint>
    struct RelaxedPow {
        RelaxedPow(long long k = 0) : k(k) {}

        mint append(const mint& fi) {
            if (k == 0) {
                return g.emplace_back(g.empty() ? 1 : 0);
            }
            static inv_mods<mint> invs;
            if (is_zero) {
                if (fi == 0) {
                    z = std::min(z + k, 1000000000LL);
                } else {
                    is_zero = false;
                    inv_base = fi.inv();
                }
            }
            if (not is_zero) {
                f.push_back(fi);
            }
            if (index < z) {
                g.push_back(0);
            } else if (index == z) {
                g.push_back(f[0].pow(k));
            } else {
                int i = index - z;
                mint v1 = fg1.append(mint(k - (i - 1)) * g[z + i - 1], f[i]);
                mint v2 = fg2.append(g[z + i - 1], mint(k) * (i - 1) * f[i]);
                g.push_back((v1 + v2) * inv_base * invs[i]);
            }
            ++index;
            return g.back();
        }
        mint operator[](int i) const {
            return g[i];
        }
    private:
        long long k;
        long long z = 0;
        long long index = 0;
        bool is_zero = true;
        mint inv_base = 0;

        std::vector<mint> f, g;
        RelaxedConvolutionNTT<mint> fg1;
        RelaxedConvolutionNTT<mint> fg2;
    };

    template <typename mint>
    struct RelaxedSqrt {
        std::optional<mint> append(const mint& fi) {
            if (g.empty()) {
                auto opt_g0 = safe_sqrt(fi);
                if (not opt_g0) return std::nullopt;
                mint g0 = *opt_g0;
                c = (2 * g0).inv();
                return g.emplace_back(g0);
            } else if (g.size() == 1) {
                return g.emplace_back(c * fi);
            } else {
                mint gi = c * (fi - gg.append(g.back(), g.back()));
                return g.emplace_back(gi);
            }
        }
        mint operator[](int i) const {
            return g[i];
        }
    private:
        mint c = 0;
        std::vector<mint> g;
        RelaxedConvolutionNTT<mint> gg;
    };
} // namespace suisen

Formal Power Series Relaxed (library/polynomial/formal_power_series_relaxed.hpp)