#921861 (C++17(clang)) No.215 素数サイコロと合成数サイコロ (3-Hard)

提出ソース
結果

問題	No.215 素数サイコロと合成数サイコロ (3-Hard)
ユーザー	SPARKLE_math
提出日時	2023-10-22 18:46:39
言語	C++17(clang) (17.0.6 + boost 1.83.0)
結果	AC
実行時間	950 ms / 4,000 ms
コード長	17,057 bytes
コンパイル時間	2,062 ms
コンパイル使用メモリ	154,536 KB
実行使用メモリ	23,180 KB
最終ジャッジ日時	2024-09-22 09:50:10
合計ジャッジ時間	5,377 ms
ジャッジサーバーID （参考情報）	judge3 / judge4
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	921 ms 23,180 KB
testcase_01	AC	950 ms 23,140 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
#line 1 "main.cpp"
#include <iostream>
#include <algorithm>
#include <vector>
#include <iomanip>
#include <math.h>
#include <functional>
#include <map>
#include <set>
#include <string>
#include <unordered_map>
#include <unordered_set>
#include <queue>
#include <stack>
#include <cstring>
#include <assert.h>
#include <sys/mman.h>
#include <unistd.h>
#include <chrono>
#include <numeric>
#include <cstdint>
#line 2 "/home/kokoro601/compro_library/Utils/FastIO.hpp"
#include <string.h>
#line 4 "/home/kokoro601/compro_library/Utils/FastIO.hpp"


namespace fastio{
    static constexpr size_t buf_size = 1 << 18;
    static constexpr size_t integer_size = 20;
    static constexpr size_t block_size = 10000;

    static char inbuf[buf_size + 1] = {};
    static char outbuf[buf_size + 1] = {};
    static char block_str[block_size * 4 + 1] = {};

    static constexpr uint64_t power10[] = {
        1, 10, 100, 1000, 10000, 100000, 1000000, 10000000, 100000000,
        1000000000, 10000000000, 100000000000, 1000000000000, 10000000000000,
        100000000000000, 1000000000000000, 10000000000000000, 100000000000000000,
        1000000000000000000, 10000000000000000000u
    };

    struct Scanner {
        private:
        size_t pos,end;

        void load() {
            end = fread(inbuf,1,buf_size,stdin);
            inbuf[end] = '\0';
        }
        void reload() {
            size_t len = end - pos;
            memmove(inbuf,inbuf + pos,len);
            end = len + fread(inbuf + len,1,buf_size - len,stdin);
            inbuf[end] = '\0';
            pos = 0;
        }
        void skip_space() {
            while(inbuf[pos] <= ' '){
                if(__builtin_expect(++pos == end, 0)) reload();
            }
        }
        char get_next() { return inbuf[pos++]; }
        char get_next_nonspace() {
            skip_space();
            return inbuf[pos++];
        }
        public:
        Scanner() { load(); }

        void scan(char& c) { c = get_next_nonspace(); }
        void scan(std::string& s){
            skip_space();
            s = "";
            do {
                size_t start = pos;
                while (inbuf[pos] > ' ') pos++;
                s += std::string(inbuf + start, inbuf + pos);
                if (inbuf[pos] !='\0') break;
                reload();
            } while (true);
        }

        template <class T>
        typename std::enable_if<std::is_integral<T>::value, void>::type scan(T &x) {
            char c = get_next_nonspace();
            if(__builtin_expect(pos + integer_size >= end, 0)) reload();
            bool neg = false;
            if (c == '-') neg = true, x = 0;
            else x = c & 15;
            while ((c = get_next()) >= '0') x = x * 10 + (c & 15);
            if (neg) x = -x;
        }

        template <class Head, class... Others>
        void scan(Head& head, Others&... others) {
            scan(head); scan(others...);
        }

        template <class T>
        Scanner& operator >> (T& x) {
            scan(x); return *this;
        }
    };

    struct Printer {
        private:
        size_t pos = 0;
        
        void flush() {
            fwrite(outbuf, 1, pos, stdout);
            pos = 0;
        }

        void pre_calc() {
            for (size_t i = 0; i < block_size; i++) {
                size_t j = 4, k = i;
                while (j--) {
                    block_str[i * 4 + j] = k % 10 + '0';
                    k /= 10;
                }
            }
        }

        static constexpr size_t get_integer_size(uint64_t n) {
            if(n >= power10[10]) {
                if (n >= power10[19]) return 20;
                if (n >= power10[18]) return 19;
                if (n >= power10[17]) return 18;
                if (n >= power10[16]) return 17;
                if (n >= power10[15]) return 16;
                if (n >= power10[14]) return 15;
                if (n >= power10[13]) return 14;
                if (n >= power10[12]) return 13;
                if (n >= power10[11]) return 12;
                return 11;
            }
            else {
                if (n >= power10[9]) return 10;
                if (n >= power10[8]) return 9;
                if (n >= power10[7]) return 8;
                if (n >= power10[6]) return 7;
                if (n >= power10[5]) return 6;
                if (n >= power10[4]) return 5;
                if (n >= power10[3]) return 4;
                if (n >= power10[2]) return 3;
                if (n >= power10[1]) return 2;
                return 1;
            }
        }

        public:
        Printer() { pre_calc(); }
        ~Printer() { flush(); }

        void print(char c){
            outbuf[pos++] = c;
            if (__builtin_expect(pos == buf_size, 0)) flush();
        }
        void print(const char *s) {
            while(*s != 0) {
                outbuf[pos++] = *s++;
                // if (pos == buf_size) flush();
                if (__builtin_expect(pos == buf_size, 0)) flush();
            }
        }
        void print(const std::string& s) {
            for(auto c : s){
                outbuf[pos++] = c;
                // if (pos == buf_size) flush();
                if (__builtin_expect(pos == buf_size, 0)) flush();
            }
        }

        template <class T>
        typename std::enable_if<std::is_integral<T>::value, void>::type print(T x) {
            if (__builtin_expect(pos + integer_size >= buf_size, 0)) flush();
            if (x < 0) print('-'), x = -x;
            size_t digit = get_integer_size(x);
            size_t len = digit;
            while (len >= 4) {
                len -= 4;
                memcpy(outbuf + pos + len, block_str + (x % block_size) * 4, 4);
                x /= block_size;
            }
            memcpy(outbuf + pos, block_str + x * 4 + (4 - len), len);
            pos += digit;
        }

        template <class Head, class... Others>
        void print(const Head& head, const Others&... others){
            print(head); print(' '); print(others...);
        }

        template <class... Args>
        void println(const Args&... args) {
            print(args...); print('\n');
        }
        
        template <class T>
        Printer& operator << (const T& x) {
            print(x); return *this;
        }
    };
};


fastio::Scanner fin;
fastio::Printer fout;
#define cin fin
#define cout fout
#line 2 "/home/kokoro601/compro_library/Math/ModComb.hpp"
// verified at:
// https://atcoder.jp/contests/abc145/submissions/33704571

template <long long MOD, int NMAX> class ModComb {
    using ll = long long;


public:
    // fac[i]: i! % MOD
    // finv[i]: i^{-1} 
    ll fac[NMAX + 1], finv[NMAX + 1], inv[NMAX + 1];
    ModComb() {
        fac[0] = fac[1] = 1;
        finv[0] = finv[1] = 1;
        inv[1] = 1;

        for (int i = 2; i < NMAX + 1; i++) {
            fac[i] = fac[i - 1] * i % MOD;
            inv[i] = MOD - inv[MOD % i] * (MOD / i) % MOD; // a^{-1} = -(p / a) * (p % a)^{-1}
            finv[i] = finv[i - 1] * inv[i] % MOD;
        }
    }

    ll calc(int n, int k) {
        if (n < k) return 0;
        if (n < 0 || k < 0) return 0;
        return fac[n] * (finv[k] * finv[n - k] % MOD) % MOD;
    }
};
#line 4 "/home/kokoro601/compro_library/Utils/ModInt.hpp"

template <uint32_t MD> struct ModInt {
    using M = ModInt;
    using uint = uint32_t;
    using ull = uint64_t;
    using ll = int64_t;
    uint v;
    ModInt(ll _v = 0) { set_v(_v % MD + MD); }
    M& set_v(uint _v) {
        v = (_v < MD) ? _v : _v - MD;
        return *this;
    }
    explicit operator bool() const { return v != 0; }
    M operator-() const { return M() - *this; }
    M operator+(const M& r) const { return M().set_v(v + r.v); }
    M operator-(const M& r) const { return M().set_v(v + MD - r.v); } // "v + MD - r.v" can exceed MD, so set_v is needed
    M operator*(const M& r) const { return M().set_v(ull(v) * r.v % MD); }
    M operator/(const M& r) const { return *this * r.inv(); }
    M& operator+=(const M& r) { return *this = *this + r; }
    M& operator-=(const M& r) { return *this = *this - r; }
    M& operator*=(const M& r) { return *this = *this * r; }
    M& operator/=(const M& r) { return *this = *this / r; }
    bool operator==(const M& r) const { return v == r.v; }
    bool operator!=(const M& r) const { return v != r.v;}
    M pow(ull n) const {
        M x = *this, r = 1;
        while (n) {
            if (n & 1) r *= x;
            x *= x;
            n >>= 1;
        }
        return r;
    }
    M inv() const { return pow(MD - 2); }
    friend std::ostream& operator<<(std::ostream& os, const M& r) { return os << r.v; }
};
#line 6 "/home/kokoro601/compro_library/DataStructure/NTT.hpp"

template<uint32_t MD, uint32_t root> class NTT {
public:
    using Mint = ModInt<MD>;
    NTT(){}

    // type : ift or not
    // a.size() should be less than 1 << 23
    void nft(bool type, std::vector<Mint>& a) {
        int n = (int)a.size(), s = 0;
        while ((1 << s) < n) s++;
        assert(1 << s == n);

        // these are calculated only once because the type is static
        static std::vector<Mint> ep, iep;
        Mint g = root;
        while (ep.size() <= s) {
            ep.push_back(g.pow(Mint(-1).v / (1 << ep.size())));
            iep.push_back(ep.back().inv());
        }


        std::vector<Mint> b(n);
        // Stockham FFT
        // no need to perform bit reversal (but not in-place)
        // memory access is sequantial
        for (int i = 1; i <= s; i++) { 
            int w = 1 << (s - i);
            Mint base = type ? iep[i] : ep[i];
            Mint now = 1;
            for (int y = 0; y < n / 2; y += w) {
                for (int x = 0; x < w; x++) {
                    auto s = a[y << 1 | x];
                    auto t = now * a[y << 1 | x | w];
                    b[y | x] = s + t;
                    b[y | x | n >> 1] = s - t;
                }
                now *= base;
            }
            swap(a, b);
        }
    }

    std::vector<Mint> convolution(const std::vector<Mint> &a, const std::vector<Mint> &b) {
        int n = (int)a.size(), m = (int)b.size();
        if (!n || !m) return {};
        int lg = 0;
        while ((1 << lg) < n + m - 1) lg++;
        int z = 1 << lg;
        // The type of a2 and b2 is std::vector<Mint> (not reference)
        // Therefore, a2 and b2 are copies of a and b, respectively.
        auto a2 = a, b2 = b;
        a2.resize(z);
        b2.resize(z);
        nft(false, a2);
        nft(false, b2);
        for (int i = 0; i < z; i++) a2[i] *= b2[i];
        nft(true, a2);
        a2.resize(n + m - 1);
        Mint iz = Mint(z).inv();
        for (int i = 0; i < n + m - 1; i++) a2[i] *= iz;
        return a2;
    }
};

#line 2 "/home/kokoro601/compro_library/Math/ExtGCD.hpp"

// Calculate the solution of ax + by = gcd(a, b)
// and return gcd(a, b)
long long extGCD(long long a, long long b, long long &x, long long &y) {
    if (b == 0) {
        x = 1;
        y = 0;
        return a;
    }

    long long d = extGCD(b, a % b, y, x);
    y -= (a / b) * x;
    return d;
}

long long gcd(long long a, long long b) {
    long long unused1, unused2;
    return extGCD(a, b, unused1, unused2);
}
#line 4 "/home/kokoro601/compro_library/Math/ModOp.hpp"

long long modpow(long long a, long long x, long long m) {
    long long ret = 1;
    while (x > 0) {
        if (x & 1) ret = ret * a % m;
        a = a * a % m;
        x >>= 1;
    }
    return ret;
}

long long modinv(long long a, long long m) {
    long long x, _;
    extGCD(a, m, x, _);
    x %= m;
    if (x < 0) x += m;
    return x;
}

#line 4 "/home/kokoro601/compro_library/Math/BostanMori.hpp"

template<int MOD>class NaiveNTT {
    using Mint = ModInt<MOD>;
public:
    NaiveNTT() {}

    std::vector<Mint> convolution(const std::vector<Mint> &a, const std::vector<Mint> &b) const {
        std::vector<Mint> ret(a.size() + b.size() - 1);
        for (int i = 0; i < a.size(); i++)
            for (int j = 0; j < b.size(); j++)
                ret[i + j] += a[i] * b[j];
        return ret;
    }
};

template<class T, class NTTType> class BostanMori{
    std::vector<T> p, q;
    NTTType ntt;
public:
    BostanMori(std::vector<T> &a, std::vector<T> &c) {
        assert(a.size() == c.size());
        int d = c.size();
        q.resize(d + 1);
        q[0] = 1;
        for (int i = 0; i < d; i++) q[i + 1] = -c[i];
        p = ntt.convolution(a, q);
        p.resize(d);
    }

    T rec(std::vector<T> _p, std::vector<T> _q, long long n) const {
        while (n) {
            assert(_q[0] == 1);
            int d = _q.size(); 
            std::vector<T> _qminus(_q);
            for (int i = 1; i < d; i += 2) _qminus[i] = -_qminus[i];
            _p = ntt.convolution(_p, _qminus);
            _q = ntt.convolution(_q, _qminus);
            for (int i = 0; i < d; i++) _q[i] = _q[2*i];
            _q.resize(d);
            if (n & 1) for (int i = 0; i < d; i++) _p[i] = _p[2*i + 1]; 
            else       for (int i = 0; i < d; i++) _p[i] = _p[2*i];
            _p.resize(d - 1);
            n >>= 1;
        }
        assert(_q[0] == 1);
        return _p[0]; 
    }

    T operator[] (const long long n) const { return rec(p, q, n); }
};
#line 27 "main.cpp"

using namespace std;
using ll = long long;
using Graph = vector<vector<int>>;
using u128 = __uint128_t;
using u64 = uint64_t;

// 番兵は入っているとする
ll garner(vector<ll> &b, vector<ll> &m) {
    vector<ll> coeffs(m.size(), 1);
    vector<ll> constants(m.size(), 0);
    for (size_t k = 0; k < b.size(); k++) {
        ll diff = (b[k] > constants[k]) ? (b[k] - constants[k]) : (b[k] - constants[k] + m[k]);
        ll t = (diff * modinv(coeffs[k], m[k])) % m[k];
        for (size_t i = k + 1; i < m.size(); i++) {
            (constants[i] += t * coeffs[i]) %= m[i];
            (coeffs[i] *= m[k]) %= m[i];
        }
    }

    return constants.back();
}

template<class T> class ArbitraryModNTT {
public:
    ArbitraryModNTT() {}

    vector<T> convolution(vector<T> &a, vector<T> &b) const {
        int n = a.size();
        int m = b.size();
        constexpr size_t MOD1 = 167772161;
        constexpr size_t MOD2 = 469762049;
        constexpr size_t MOD3 = 1224736769;
        using Mint1 = ModInt<MOD1>;
        using Mint2 = ModInt<MOD2>;
        using Mint3 = ModInt<MOD3>;

        NTT<MOD1, 3> ntt1;
        NTT<MOD2, 3> ntt2;
        NTT<MOD3, 3> ntt3;

        vector<Mint1> a1(n), b1(m);
        vector<Mint2> a2(n), b2(m);
        vector<Mint3> a3(n), b3(m);

        for (int i = 0; i < n; i++) a1[i] = (a[i].v) % MOD1, a2[i] = (a[i].v) % MOD2, a3[i] = (a[i].v) % MOD3;
        for (int i = 0; i < m; i++) b1[i] = (b[i].v) % MOD1, b2[i] = (b[i].v) % MOD2, b3[i] = (b[i].v) % MOD3;
        auto c1 = ntt1.convolution(a1, b1);
        auto c2 = ntt2.convolution(a2, b2);
        auto c3 = ntt3.convolution(a3, b3);

        vector<T> c(c1.size());
        const Mint2 m1_inv_m2 = modinv(MOD1, MOD2);
        const Mint3 m1m2_inv_m3 = modinv((Mint3(MOD1) * MOD2).v, MOD3);
        for (int i = 0; i < c1.size(); i++) {
            ll t1 = (m1_inv_m2 * ((ll)c2[i].v - c1[i].v)).v;
            ll t = (m1m2_inv_m3 * ((ll)c3[i].v - t1 * MOD1 - c1[i].v)).v;
            c[i] = T(t) * MOD1 * MOD2 + T(t1) * MOD1 + c1[i].v;
        }

        return c;
    }
};

constexpr size_t MOD = 1e9 + 7;
using Mint = ModInt<MOD>;



int prime[]     = {0, 2, 3, 5, 7, 11, 13};
int composite[] = {0, 4, 6, 8, 9, 10, 12};

Mint dp0[2][305][4000], dp1[2][305][4000];

vector<Mint> multipoly(vector<Mint> &a, vector<Mint> &b) {
    vector<Mint> c(a.size() + b.size() - 1);
    for (int i = 0; i < a.size(); i++)
        for (int j = 0; j < b.size(); j++)
            c[i + j] += a[i] * b[j];
    return c;
}

int main() {
    ll n; cin >> n;
    int p, c; cin >> p >> c;

    
    dp0[0][0][0] = 1;
    for (int i = 0; i < 6; i++) {
        for (int cnt = 0; cnt <= p; cnt++) 
            for (int sm = 0; sm <= prime[i] * cnt; sm++) 
                for (int k = 0; k <= p - cnt; k++)
                    dp0[1][cnt + k][sm + prime[i+1]*k] += dp0[0][cnt][sm];
        swap(dp0[0], dp0[1]);
        memset(dp0[1], 0, sizeof(Mint) * 305 * 4000);
    }

    dp1[0][0][0] = 1;
    for (int i = 0; i < 6; i++) {
        for (int cnt = 0; cnt <= c; cnt++)
            for (int sm = 0; sm <= composite[i] * cnt; sm++)
                for (int k = 0; k <= c - cnt; k++)
                    dp1[1][cnt + k][sm + composite[i+1]*k] += dp1[0][cnt][sm];
        swap(dp1[0], dp1[1]);
        memset(dp1[1], 0, sizeof(Mint) * 305 * 4000);
    }

    vector<Mint> v0(13*p + 1), v1(12*c + 1);
    for (int i = 0; i <= 13*p; i++) v0[i] = dp0[0][p][i];
    for (int i = 0; i <= 12*c; i++) v1[i] = dp1[0][c][i];


    vector<Mint> coeff = multipoly(v0, v1);
    coeff.erase(coeff.begin());
    int mx = 13*p + 12*c;
    vector<Mint> beg(mx, 1);


    BostanMori<Mint, ArbitraryModNTT<Mint>> bm(beg, coeff);
    cout << bm[mx + n - 1].v << "\n";
    return 0;
}
yukicoder

結果

テストケース

ソースコード
