cp-library-cpp

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

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

:warning: library/sequence/sum_of_powers.hpp

Depends on

Code

#ifndef SUISEN_SUM_POWERS
#define SUISEN_SUM_POWERS

#include "library/math/pow_mods.hpp"
#include "library/sequence/bernoulli_number.hpp"

namespace suisen {
    // res[p] = 1^p + 2^p + ... + n^p for p=0, ..., max_exponent (O(k log k), where k=max_exponent)
    template <typename FPSType>
    auto sum_of_powers(int n, int max_exponent, const std::vector<typename FPSType::value_type> &bernoulli_table) {
        const int k = max_exponent;
        assert(bernoulli_table.size() >= size_t(k + 2));
        using fps = FPSType;
        using mint = typename FPSType::value_type;

        factorial<mint> fac(k + 1);
        pow_mods<mint> pow_n(n, k + 1);

        fps f(k + 2);
        for (int j = 0; j <= k + 1; ++j) {
            f[j] = pow_n[j] * fac.fac_inv(j);
        }
        std::vector<mint> b(bernoulli_table.begin(), bernoulli_table.begin() + (k + 2));
        b[1] *= -1; // b[1] = +1/2
        for (int j = 0; j <= k + 1; ++j) {
            b[j] *= fac.fac_inv(j);
        }
        f *= b;

        std::vector<mint> res(k + 1);
        for (int p = 0; p <= k; ++p) {
            res[p] = fac.fac(p) * (f[p + 1] - b[p + 1]);
        }
        return res;
    }

    // res[p] = 1^p + 2^p + ... + n^p for p=0, ..., max_exponent (O(k log k), where k=max_exponent)
    template <typename FPSType>
    auto sum_of_powers(int n, int max_exponent) {
        return sum_of_powers<FPSType>(n, max_exponent, bernoulli_number<FPSType>(max_exponent + 1));
    }
} // namespace suisen

#endif // SUISEN_SUM_POWERS
#line 1 "library/sequence/sum_of_powers.hpp"



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



#include <vector>

namespace suisen {
    template <int base_as_int, typename mint>
    struct static_pow_mods {
        static_pow_mods() = default;
        static_pow_mods(int n) { ensure(n); }
        const mint& operator[](int i) const {
            ensure(i);
            return pows[i];
        }
        static void ensure(int n) {
            int sz = pows.size();
            if (sz > n) return;
            pows.resize(n + 1);
            for (int i = sz; i <= n; ++i) pows[i] = base * pows[i - 1];
        }
    private:
        static inline std::vector<mint> pows { 1 };
        static inline mint base = base_as_int;
        static constexpr int mod = mint::mod();
    };

    template <typename mint>
    struct pow_mods {
        pow_mods() = default;
        pow_mods(mint base, int n) : base(base) { ensure(n); }
        const mint& operator[](int i) const {
            ensure(i);
            return pows[i];
        }
        void ensure(int n) const {
            int sz = pows.size();
            if (sz > n) return;
            pows.resize(n + 1);
            for (int i = sz; i <= n; ++i) pows[i] = base * pows[i - 1];
        }
    private:
        mutable std::vector<mint> pows { 1 };
        mint base;
        static constexpr int mod = mint::mod();
    };
}


#line 1 "library/sequence/bernoulli_number.hpp"



#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];
        }
        // i の逆数
        // i = 0 の場合は assert 違反となる
        U inv(const int i) {
            assert(i > 0);
            ensure(i);
            return _fac_inv[i] * _fac[i - 1];
        }
        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];
        }
        // binom(n, r) の逆数
        // binom(n, r) = 0 の場合は assert 違反となる
        U binom_inv(const int n, const int r) {
            assert(r >= 0 and n >= r);
            ensure(n);
            return _fac_inv[n] * _fac[r] * _fac[n - r];
        }
        // n 種類から重複を許して r 個選ぶ場合の数
        // x_1+x_2+...+x_n=r（x_i は非負整数）となる x の個数でもある
        // multichoose(n, r) = binom(n + r - 1, r)
        U multichoose(const int n, const int r) {
            if (n < 0 or r < 0) return 0;
            return r > 0 ? binom(n + r - 1, r) : U(1);
        }
        // n 種類から重複を許して r 個選ぶ場合の数 multichoose(n, r) の逆数
        // x_1+x_2+...+x_n=r（x_i は非負整数）となる x の個数の逆数でもある
        // multichoose(n, r) = binom(n + r - 1, r)
        // multichoose(n, r) = 0 の場合は assert 違反となる
        U multichoose_inv(const int n, const int r) {
            assert(n >= 0 and r >= 0);
            return r > 0 ? binom_inv(n + r - 1, r) : U(1);
        }
        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];
        }
        // perm(n, r) の逆数
        // perm(n, r) = 0 の場合は assert 違反となる
        U perm_inv(const int n, const int r) {
            assert(r >= 0 and n >= r);
            ensure(n);
            return _fac_inv[n] * _fac[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 5 "library/sequence/bernoulli_number.hpp"

namespace suisen {
    template <typename FPSType>
    std::vector<typename FPSType::value_type> bernoulli_number(int n) {
        using mint = typename FPSType::value_type;
        factorial<mint> fac(n);
        FPSType a(n + 1);
        for (int i = 0; i <= n; ++i) a[i] = fac.fac_inv(i + 1);
        a.inv_inplace(n + 1), a.resize(n + 1);
        for (int i = 2; i <= n; ++i) a[i] *= fac(i);
        return a;
    }
} // namespace suisen


#line 6 "library/sequence/sum_of_powers.hpp"

namespace suisen {
    // res[p] = 1^p + 2^p + ... + n^p for p=0, ..., max_exponent (O(k log k), where k=max_exponent)
    template <typename FPSType>
    auto sum_of_powers(int n, int max_exponent, const std::vector<typename FPSType::value_type> &bernoulli_table) {
        const int k = max_exponent;
        assert(bernoulli_table.size() >= size_t(k + 2));
        using fps = FPSType;
        using mint = typename FPSType::value_type;

        factorial<mint> fac(k + 1);
        pow_mods<mint> pow_n(n, k + 1);

        fps f(k + 2);
        for (int j = 0; j <= k + 1; ++j) {
            f[j] = pow_n[j] * fac.fac_inv(j);
        }
        std::vector<mint> b(bernoulli_table.begin(), bernoulli_table.begin() + (k + 2));
        b[1] *= -1; // b[1] = +1/2
        for (int j = 0; j <= k + 1; ++j) {
            b[j] *= fac.fac_inv(j);
        }
        f *= b;

        std::vector<mint> res(k + 1);
        for (int p = 0; p <= k; ++p) {
            res[p] = fac.fac(p) * (f[p + 1] - b[p + 1]);
        }
        return res;
    }

    // res[p] = 1^p + 2^p + ... + n^p for p=0, ..., max_exponent (O(k log k), where k=max_exponent)
    template <typename FPSType>
    auto sum_of_powers(int n, int max_exponent) {
        return sum_of_powers<FPSType>(n, max_exponent, bernoulli_number<FPSType>(max_exponent + 1));
    }
} // namespace suisen
Back to top page