けんちょんの競プロ精進記録

競プロの精進記録や小ネタを書いていきます

Yosupo Library Checker - Partition Function

ここでは  O(N\sqrt{N}) の計算量の解法を実装した。本当は FPS を使えば  O(N \log N) の計算量でできる。

問題概要

非負整数値  N が与えられる。

 0, 1, \dots, N の分割数を 998244353 で割った余りを求めよ。

制約

  •  0 \le N \le 5 \times 10^{5}

考えたこと

ここでは、この記事に書いた  O(N\sqrt{N}) の計算量の解法を実装した。

qiita.com

具体的には Euler の五角数定理によって導出される、次の公式がある。

別解

FPS を用いた  O(N \log N) 解法を Shirotsume さんが書いている。とてもわかりやすい。

ladywingclover.hatenablog.com

コード

#include <bits/stdc++.h>
using namespace std;


// modint
template<int MOD> struct Fp {
    // inner value
    long long val;
    
    // constructor
    constexpr Fp() noexcept : val(0) { }
    constexpr Fp(long long v) noexcept : val(v % MOD) {
        if (val < 0) val += MOD;
    }
    constexpr long long get() const noexcept { return val; }
    constexpr int get_mod() const noexcept { return MOD; }
    
    // arithmetic operators
    constexpr Fp operator - () const noexcept {
        return val ? MOD - val : 0;
    }
    constexpr Fp operator + (const Fp &r) const noexcept { return Fp(*this) += r; }
    constexpr Fp operator - (const Fp &r) const noexcept { return Fp(*this) -= r; }
    constexpr Fp operator * (const Fp &r) const noexcept { return Fp(*this) *= r; }
    constexpr Fp operator / (const Fp &r) const noexcept { return Fp(*this) /= r; }
    constexpr Fp& operator += (const Fp &r) noexcept {
        val += r.val;
        if (val >= MOD) val -= MOD;
        return *this;
    }
    constexpr Fp& operator -= (const Fp &r) noexcept {
        val -= r.val;
        if (val < 0) val += MOD;
        return *this;
    }
    constexpr Fp& operator *= (const Fp &r) noexcept {
        val = val * r.val % MOD;
        return *this;
    }
    constexpr Fp& operator /= (const Fp &r) noexcept {
        long long a = r.val, b = MOD, u = 1, v = 0;
        while (b) {
            long long t = a / b;
            a -= t * b, swap(a, b);
            u -= t * v, swap(u, v);
        }
        val = val * u % MOD;
        if (val < 0) val += MOD;
        return *this;
    }
    constexpr Fp pow(long long n) const noexcept {
        Fp res(1), mul(*this);
        while (n > 0) {
            if (n & 1) res *= mul;
            mul *= mul;
            n >>= 1;
        }
        return res;
    }
    constexpr Fp inv() const noexcept {
        Fp res(1), div(*this);
        return res / div;
    }

    // other operators
    constexpr bool operator == (const Fp &r) const noexcept {
        return this->val == r.val;
    }
    constexpr bool operator != (const Fp &r) const noexcept {
        return this->val != r.val;
    }
    friend constexpr istream& operator >> (istream &is, Fp<MOD> &x) noexcept {
        is >> x.val;
        x.val %= MOD;
        if (x.val < 0) x.val += MOD;
        return is;
    }
    friend constexpr ostream& operator << (ostream &os, const Fp<MOD> &x) noexcept {
        return os << x.val;
    }
    friend constexpr Fp<MOD> modpow(const Fp<MOD> &r, long long n) noexcept {
        return r.pow(n);
    }
    friend constexpr Fp<MOD> modinv(const Fp<MOD> &r) noexcept {
        return r.inv();
    }
};

// partition number
template<class T> struct PartitionFunction {
    vector<T> P;
    constexpr PartitionFunction(int MAX) noexcept : P(MAX+1, 0) {
        P[0] = 1;
        for (int i = 1; i <= MAX; ++i) {
            for (int j = 1, sign = 1; i-(j*j*3-j)/2 >= 0; ++j, sign *= -1) {
                P[i] += P[i-(j*j*3-j)/2] * sign;
                if (i-(j*j*3+j)/2 >= 0) P[i] += P[i-(j*j*3+j)/2] * sign;
            }
        }
    }
    constexpr T get(int n) {
        if (n < 0) return 0;
        return P[n];
    }
};

void Yosupo_Partition_Function() {
    using mint = Fp<998244353>;
    int N;
    cin >> N;
    PartitionFunction<mint> pfun(N);
    for (int i = 0; i <= N; ++i) cout << pfun.get(i) << " ";
    cout << endl;
}

int main() {
    Yosupo_Partition_Function();
}