#1061250 (C++23) No.1145 Sums of Powers

• 問題: No.1145 Sums of Powers
• ユーザー: SnowBeenDiding
• 提出日時: 2025-03-30 11:48:57
• 言語: C++23 (gcc 13.3.0 + boost 1.87.0)
• 結果: AC
• 実行時間: 378 ms / 2,000 ms
• コード長: 5,027 bytes
• コンパイル時間: 7,095 ms
• コンパイル使用メモリ: 333,360 KB
• 実行使用メモリ: 42,228 KB
• 最終ジャッジ日時: 2025-03-30 11:49:07
• 合計ジャッジ時間: 8,749 ms
• ジャッジサーバーID （参考情報）: judge2 / judge4

ファイルパターン	結果
other	AC * 6

ソースコード

#include <atcoder/all>
#include <bits/stdc++.h>
#define rep(i, a, b) for (ll i = (ll)(a); i < (ll)(b); i++)
using namespace atcoder;
using namespace std;

typedef long long ll;

struct PowerSum {
    template <class T> struct FormalPowerSeries : vector<T> {
        using vector<T>::vector;
        using vector<T>::operator=;
        using F = FormalPowerSeries;

        F operator-() const {
            F res(*this);
            for (auto &e : res)
                e = -e;
            return res;
        }
        F &operator*=(const T &g) {
            for (auto &e : *this)
                e *= g;
            return *this;
        }
        F &operator/=(const T &g) {
            assert(g != T(0));
            *this *= g.inv();
            return *this;
        }
        F &operator+=(const F &g) {
            int n = (*this).size(), m = g.size();
            for (int i = 0; i < min(n, m); i++) {
                (*this)[i] += g[i];
            }
            return *this;
        }
        F &operator-=(const F &g) {
            int n = (*this).size(), m = g.size();
            for (int i = 0; i < min(n, m); i++) {
                (*this)[i] -= g[i];
            }
            return *this;
        }
        F inv(int d = -1) const {
            int n = (*this).size();
            assert(n != 0 && (*this)[0] != 0);
            if (d == -1)
                d = n;
            assert(d > 0);
            F res{(*this)[0].inv()};
            while (res.size() < d) {
                int m = size(res);
                F f(begin(*this), begin(*this) + min(n, 2 * m));
                F r(res);
                f.resize(2 * m), internal::butterfly(f);
                r.resize(2 * m), internal::butterfly(r);
                for (int i = 0; i < 2 * m; i++) {
                    f[i] *= r[i];
                }
                internal::butterfly_inv(f);
                f.erase(f.begin(), f.begin() + m);
                f.resize(2 * m), internal::butterfly(f);
                for (int i = 0; i < 2 * m; i++) {
                    f[i] *= r[i];
                }
                internal::butterfly_inv(f);
                T iz = T(2 * m).inv();
                iz *= -iz;
                for (int i = 0; i < m; i++) {
                    f[i] *= iz;
                }
                res.insert(res.end(), f.begin(), f.begin() + m);
            }
            return {res.begin(), res.begin() + d};
        }

        // fast: FMT-friendly modulus only
        F &operator*=(const F &g) {
            int n = (*this).size();
            *this = convolution(*this, g);
            return *this;
        }
        F &operator/=(const F &g) {
            int n = (*this).size();
            *this = convolution(*this, g.inv(n));
            return *this;
        }

        F operator*(const T &g) const { return F(*this) *= g; }
        F operator/(const T &g) const { return F(*this) /= g; }
        F operator+(const F &g) const { return F(*this) += g; }
        F operator-(const F &g) const { return F(*this) -= g; }
        F operator<<(const int d) const { return F(*this) <<= d; }
        F operator>>(const int d) const { return F(*this) >>= d; }
        F operator*(const F &g) const { return F(*this) *= g; }
        F operator/(const F &g) const { return F(*this) /= g; }
        F operator*(vector<pair<int, T>> g) const { return F(*this) *= g; }
        F operator/(vector<pair<int, T>> g) const { return F(*this) /= g; }
    };

    using fps = FormalPowerSeries<modint998244353>;

    int mod;
    vector<int> a;

    PowerSum() {}
    PowerSum(vector<int> a, int mod) : a(a), mod(mod) {
        modint::set_mod(mod);
        for (auto &x : a)
            x %= mod;
    }

    /**
     * S[k] = Σ[i=1,n] A[i]^k
     * O( N(log(N))^2 + Mlog(M) )
     * https://yukicoder.me/problems/no/1145
     */
    vector<int> sums_of_powers(vector<int> a, int m) {
        assert(mod == 998244353);
        int n = a.size();
        vector<pair<fps, fps>> f(n * 2);
        for (int i = n; i < n * 2; i++) {
            f[i].first = fps{1};
            f[i].second = fps{1, -a[i - n]};
        }
        for (int i = n - 1; i > 0; i--) {
            f[i].first = f[i * 2].first * f[i * 2 + 1].second +
                         f[i * 2 + 1].first * f[i * 2].second;
            f[i].second = f[i * 2].second * f[i * 2 + 1].second;
        }
        while (f[1].first.size() < m) {
            f[1].first.resize(m);
            f[1].second.resize(m);
        }
        vector<int> ret(m);
        auto v = f[1].first * f[1].second.inv();
        for (int i = 0; i < m; i++) {
            ret[i] = v[i].val();
        }
        return ret;
    }
};

int main() {
    cin.tie(0);
    cout.tie(0);
    ios::sync_with_stdio(0);
    int n, m;
    cin >> n >> m;
    vector<int> a(n);
    for (int i = 0; i < n; i++) {
        cin >> a[i];
    }
    PowerSum ps(a, 998244353);
    auto res = ps.sums_of_powers(a, m + 1);
    for (int i = 1; i < m + 1; i++) {
        cout << res[i] << " ";
    }
    cout << endl;
}