#945200 (C++17) No.2613 Sum of Combination

提出ソース
結果

問題	No.2613 Sum of Combination
コンテスト
ユーザー	Nachia
提出日時	2024-01-19 22:56:16
言語	C++17 (gcc 13.3.0 + boost 1.87.0)
結果	AC
実行時間	174 ms / 4,500 ms
コード長	14,170 bytes
コンパイル時間	2,962 ms
コンパイル使用メモリ	155,096 KB
最終ジャッジ日時	2025-02-18 21:29:05
ジャッジサーバーID （参考情報）	judge2 / judge1
このコードへのチャレンジ
（要ログイン）
ファイルパターン	結果
sample	AC * 3
other	AC * 49
権限があれば一括ダウンロードができます
ソースコード

raw source code
#ifdef NACHIA
// #define _GLIBCXX_DEBUG
#else
#define NDEBUG
#endif
#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
#include <utility>
#include <queue>
#include <array>
#include <cmath>
#include <atcoder/modint>
#include <atcoder/convolution>
using namespace std;
using i64 = long long;
using u64 = unsigned long long;
#define rep(i,n) for(i64 i=0; i<(i64)(n); i++)
#define repr(i,n) for(i64 i=(i64)(n)-1; i>=0; i--)
const i64 INF = 1001001001001001001;
const char* yn(bool x){ return x ? "Yes" : "No"; }
template<typename A> void chmin(A& l, const A& r){ if(r < l) l = r; }
template<typename A> void chmax(A& l, const A& r){ if(l < r) l = r; }
template<typename A> using nega_queue = priority_queue<A,vector<A>,greater<A>>;
using Modint = atcoder::static_modint<998244353>;

#include <iterator>
#include <functional>

template<class Elem> struct vec;

template<class Iter>
struct seq_view{
    using Ref = typename std::iterator_traits<Iter>::reference;
    using Elem = typename std::iterator_traits<Iter>::value_type;
    Iter a, b;
    Iter begin() const { return a; }
    Iter end() const { return b; }
    int size() const { return (int)(b-a); }
    seq_view(Iter first, Iter last) : a(first), b(last) {}
    seq_view sort() const { std::sort(a, b); return *this; }
    Ref& operator[](int x){ return *(a+x); }
    template<class F = std::less<Elem>, class ret = vec<int>> ret sorti(F f = F()) const {
        ret x(size()); for(int i=0; i<size(); i++) x[i] = i;
        x().sort([&](int l, int r){ return f(a[l],a[r]); });
        return x;
    }
    template<class ret = vec<Elem>> ret col() const { return ret(begin(), end()); }
    template<class F = std::equal_to<Elem>, class ret = vec<std::pair<Elem, int>>>
    ret rle(F eq = F()) const {
        auto x = ret();
        for(auto& a : (*this)){
            if(x.size() == 0 || !eq(x[x.size()-1].first, a)) x.emp(a, 1); else x[x.size()-1].second++;
        } return x;
    }
    template<class F> seq_view sort(F f) const { std::sort(a, b, f); return *this; }
    Iter uni() const { return std::unique(a, b); }
    Iter lb(const Elem& x) const { return std::lower_bound(a, b, x); }
    Iter ub(const Elem& x) const { return std::upper_bound(a, b, x); }
    int lbi(const Elem& x) const { return lb(x) - a; }
    int ubi(const Elem& x) const { return ub(x) - a; }
    seq_view bound(const Elem& l, const Elem& r) const { return { lb(l), lb(r) }; }
    template<class F> Iter lb(const Elem& x, F f) const { return std::lower_bound(a, b, x, f); }
    template<class F> Iter ub(const Elem& x, F f) const { return std::upper_bound(a, b, x, f); }
    template<class F> Iter when_true_to_false(F f) const {
        if(a == b) return a;
        return std::lower_bound(a, b, *a,
            [&](const Elem& x, const Elem&){ return f(x); });
    }
    seq_view same(Elem x) const { return { lb(x), ub(x) }; }
    template<class F> auto map(F f) const {
        vec<typename Iter::value_type> r;
        for(auto& x : *this) r.emp(f(x));
        return r;
    }
    Iter max() const { return std::max_element(a, b); }
    Iter min() const { return std::min_element(a, b); }
    template<class F = std::less<Elem>>
    Iter min(F f) const { return std::min_element(a, b, f); }
    seq_view rev() const { std::reverse(a, b); return *this; }
};

template<class Elem>
struct vec {
    using Base = typename std::vector<Elem>;
    using Iter = typename Base::iterator;
    using CIter = typename Base::const_iterator;
    using View = seq_view<Iter>;
    using CView = seq_view<CIter>;

    vec(){}
    explicit vec(int n, const Elem& value = Elem()) : a(0<n?n:0, value) {}
    template <class I2> vec(I2 first, I2 last) : a(first, last) {}
    vec(std::initializer_list<Elem> il) : a(std::move(il)) {}
    vec(Base b) : a(std::move(b)) {}
    operator Base() const { return a; }

    Iter begin(){ return a.begin(); }
    CIter begin() const { return a.begin(); }
    Iter end(){ return a.end(); }
    CIter end() const { return a.end(); }
    int size() const { return a.size(); }
    bool empty() const { return a.empty(); }
    Elem& back(){ return a.back(); }
    const Elem& back() const { return a.back(); }
    vec sortunied(){ vec x = *this; x().sort(); x.a.erase(x().uni(), x.end()); return x; }
    Iter operator()(int x){ return a.begin() + x; }
    CIter operator()(int x) const { return a.begin() + x; }
    View operator()(int l, int r){ return { (*this)(l), (*this)(r) }; }
    CView operator()(int l, int r) const { return { (*this)(l), (*this)(r) }; }
    View operator()(){ return (*this)(0,size()); }
    CView operator()() const { return (*this)(0,size()); }
    Elem& operator[](int x){ return a[x]; }
    const Elem& operator[](int x) const { return a[x]; }
    Base& operator*(){ return a; }
    const Base& operator*() const { return a; }
    vec& push(Elem args){
        a.push_back(std::move(args));
        return *this;
    }
    template<class... Args>
    vec& emp(Args &&... args){
        a.emplace_back(std::forward<Args>(args) ...);
        return *this;
    }
    template<class Range>
    vec& app(Range& x){ for(auto& v : a) emp(v); }
    Elem pop(){
        Elem x = std::move(a.back());
        a.pop_back(); return x;
    }
    bool operator==(const vec& r) const { return a == r.a; }
    bool operator!=(const vec& r) const { return a != r.a; }
    bool operator<(const vec& r) const { return a < r.a; }
    bool operator<=(const vec& r) const { return a <= r.a; }
    bool operator>(const vec& r) const { return a > r.a; }
    bool operator>=(const vec& r) const { return a >= r.a; }
    vec<vec<Elem>> pile(int n) const { return vec<vec<Elem>>(n, *this); }
    template<class F> vec& filter(F f){
        int p = 0;
        for(int q=0; q<size(); q++) if(f(a[q])) std::swap(a[p++],a[q]);
        a.resize(p); return *this;
    }
private: Base a;
};

template<class IStr, class U, class T>
IStr& operator>>(IStr& is, vec<std::pair<U,T>>& v){ for(auto& x:v){ is >> x.first >> x.second; } return is; }
template<class IStr, class T>
IStr& operator>>(IStr& is, vec<T>& v){ for(auto& x:v){ is >> x; } return is; }
template<class OStr, class T>
OStr& operator<<(OStr& os, const vec<T>& v){
    for(int i=0; i<v.size(); i++){
        if(i){ os << ' '; } os << v[i];
    } return os;
}

namespace nachia{

template<class Modint>
class Comb{
private:
    std::vector<Modint> F;
    std::vector<Modint> iF;
public:
    void extend(int newN){
        int prevN = (int)F.size() - 1;
        if(prevN >= newN) return;
        F.resize(newN+1);
        iF.resize(newN+1);
        for(int i=prevN+1; i<=newN; i++) F[i] = F[i-1] * Modint::raw(i);
        iF[newN] = F[newN].inv();
        for(int i=newN; i>prevN; i--) iF[i-1] = iF[i] * Modint::raw(i);
    }
    Comb(int n = 1){
        F.assign(2, Modint(1));
        iF.assign(2, Modint(1));
        extend(n);
    }
    Modint factorial(int n) const { return F[n]; }
    Modint invFactorial(int n) const { return iF[n]; }
    Modint invOf(int n) const { return iF[n] * F[n-1]; }
    Modint comb(int n, int r) const {
        if(n < 0 || n < r || r < 0) return Modint(0);
        return F[n] * iF[r] * iF[n-r];
    }
    Modint invComb(int n, int r) const {
        if(n < 0 || n < r || r < 0) return Modint(0);
        return iF[n] * F[r] * F[n-r];
    }
    Modint perm(int n, int r) const {
        if(n < 0 || n < r || r < 0) return Modint(0);
        return F[n] * iF[n-r];
    }
    Modint invPerm(int n, int r) const {
        if(n < 0 || n < r || r < 0) return Modint(0);
        return iF[n] * F[n-r];
    }
    Modint operator()(int n, int r) const { return comb(n,r); }
};

} // namespace nachia
#include <initializer_list>

namespace nachia{

bool IsPrime(unsigned long long x) noexcept {
    if(x <= 1) return false;
    if(x % 2 == 0) return x == 2;
    using u64 = unsigned long long;
    using u128 = __uint128_t;
    u64 d = x-1;
    int s = 0;
    int q = 63;
    while(!(d&1)){ d >>= 1; s++; }
    while(!(d >> q)) q--;
    u64 r = x; for(int t=0; t<6; t++) r*=2-r*x;
    u128 n2 = -(u128)x % x;
    auto red = [=](u128 t) noexcept -> u64 {
        t = (t + (u128)((u64)t*-r)*x) >> 64;
        return (t >= x) ? t-x : t;
    };
    u64 one = red(n2);
    for(u64 base : { 2, 325, 9375, 28178, 450775, 9780504, 1795265022 }){
        if(base%x==0) continue;
        u64 a = base = red(base%x*n2);
        for(int e=q-1; e>=0; e--){ a = red((u128)a*a); if((d>>e)&1) a = red((u128)a*base); }
        if(a == one) continue;
        for(int t=1; t<s&&a!=x-one; t++) a = red((u128)a*a);
        if(a != x-one) return false;
    }
    return true;
}

} // namespace nachia

namespace nachia{

int Popcount(unsigned long long c) noexcept {
#ifdef __GNUC__
    return __builtin_popcountll(c);
#else
    c = (c & (~0ull/3)) + ((c >> 1) & (~0ull/3));
    c = (c & (~0ull/5)) + ((c >> 2) & (~0ull/5));
    c = (c & (~0ull/17)) + ((c >> 4) & (~0ull/17));
    c = (c * (~0ull/257)) >> 56;
    return c;
#endif
}

// please ensure x != 0
int MsbIndex(unsigned long long x) noexcept {
#ifdef __GNUC__
    return 63 - __builtin_clzll(x);
#else
    using u64 = unsigned long long;
    int q = (x >> 32) ? 32 : 0;
    auto m = x >> q;
    constexpr u64 hi = 0x8888'8888;
    constexpr u64 mi = 0x1111'1111;
    m = (((m | ~(hi - (m & ~hi))) & hi) * mi) >> 35;
    m = (((m | ~(hi - (x & ~hi))) & hi) * mi) >> 31;
    q += (m & 0xf) << 2;
    q += 0x3333'3333'2222'1100 >> (((x >> q) & 0xf) << 2) & 0xf;
    return q;
#endif
}

// please ensure x != 0
int LsbIndex(unsigned long long x) noexcept {
#ifdef __GNUC__
    return __builtin_ctzll(x);
#else
    return MsbIndex(x & -x);
#endif
}

}


namespace nachia{

std::vector<std::pair<unsigned long long, int>> Factorize(unsigned long long x){
    if(x == 1) return {};
    if(IsPrime(x)) return {{x,1}};
    using u64 = unsigned long long;
    using u128 = __uint128_t;
    u64 X = x;
    std::vector<u64> p;
    for(u64 i=2; i<100; i+=1+i%2) if(x%i==0){ p.push_back(i); while(x%i==0) x/=i; }
    u64 r=1; u128 n2=1;
    auto updX = [&](){
        r = x; for(int t=0; t<6; t++) r*=2-r*x;
        n2 = -(u128)x % x;
    };
    auto red = [&](u128 t) noexcept -> u64 {
        u64 s = ((u128)x*((u64)t*r)) >> 64;
        u64 t2 = t >> 64;
        return t2-s + (t2 < s ? x : 0);
    };
    auto mult = [&](u64 a, u64 b) noexcept { return red((u128)red((u128)a*n2)*b); };
    auto gcd = [](u64 a, u64 b) noexcept {
        if(!a || !b) return a|b;
        int q = LsbIndex(a|b);
        b >>= LsbIndex(b);
        a >>= LsbIndex(a);
        while(a!=b){
            if(a<b){ b-=a; b>>=LsbIndex(b); }
            else{ a-=b; a>>=LsbIndex(a); }
        }
        return a<<q;
    };
    static u64 v = 7001;
    p.push_back(x);
    for(int pi=p.size()-1; pi<(int)p.size(); pi++) while(p[pi] != 1 && !IsPrime(p[pi])){
        x = p[pi]; updX();
        while(p[pi] == x){
            v^=v<<13; v^=v>>7; v^=v<<17; // Xorshift https://www.jstatsoft.org/article/download/v008i14/916
            u64 c = red(v); if(c == 0) continue;
            auto f = [=](u64 a) noexcept -> u64 { return red((u128)a*a+c); };
            u64 a=0, b=f(a);
            u64 buf = 1, sz = 1, nx = 10;
            while(true){
                while(nx != sz && a != b){
                    buf = mult(buf, a<=b?b-a:a-b); sz++;
                    a = f(a); b = f(f(b));
                }
                u64 g = gcd(buf, x);
                if(g != 1){
                    while(p[pi] % g == 0) p[pi] /= g;
                    p.push_back(g);
                    break;
                }
                if(a == b) break;
                nx = sz * 3 / 2;
            }
        }
    }
    std::vector<std::pair<u64, int>> res;
    for(u64 q : p) if(q != 1){
        int e=0; while(X%q == 0){ e++; X/=q; }
        if(e) res.push_back({ q, e });
    }
    return res;
}

unsigned long long Totient(unsigned long long x){
    auto F = Factorize(x);
    for(auto f : F) x -= x / f.first;
    return x;
}

} // namespace nachia
#include <cstdint>
#include <cassert>

namespace nachia{

unsigned long long PrimitiveRootDyn(unsigned long long p){
    using u64 = unsigned long long;
    using u128 = __uint128_t;
    assert(IsPrime(p));
    if(p == 2) return 1;
    auto F = Factorize(p - 1);
    u64 r = p; for(int t=0; t<6; t++) r*=2-r*p;
    u64 n2 = -(u128)p % p;
    auto red = [&](u128 t) noexcept -> u64 {
        t = (t + (u128)((u64)t*-r)*p) >> 64;
        return (t >= p) ? t-p : t;
    };
    auto mult = [&](u64 a, u64 b) noexcept { return red((u128)red((u128)a*b)*n2); };
    auto powm = [&](u64 a, u64 i) noexcept {
        u64 b = 1;
        while(i){ if(i&1){ b=mult(a,b); } a=mult(a,a); i/=2; }
        return b;
    };
    //static u64 v = 7001;
    u64 v = 1;
    while(true){
        //v^=v<<13; v^=v>>7; v^=v<<17; // Xorshift https://www.jstatsoft.org/article/download/v008i14/916
        v++;
        u64 vv = v%p;
        if(vv == 0) continue;
        bool ok = true;
        for(auto f : F){
            u64 i = (p-1)/f.first;
            if(powm(vv,i) == 1){ ok = false; break; }
        }
        if(ok) break;
    }
    return v%p;
}

} // namespace nachia

void testcase(){
    i64 N, P; cin >> N >> P;
    auto cycomb = [&](const vector<Modint>& a, const vector<Modint>& b){
        auto res = atcoder::convolution(a, b);
        rep(i,P-2) res[i] += res[P-1+i];
        res.resize(P-1);
        return res;
    };
    using dynMod = atcoder::modint;
    dynMod::set_mod(P);
    auto comb = nachia::Comb<dynMod>(P-1);
    i64 g = nachia::PrimitiveRootDyn(P);
    vector<i64> discreteLog(P);
    for(i64 i=0, p=1; i<P-1; i++, p=p*g%P) discreteLog[p] = i;
    vector<Modint> res(P-1);
    res[0] = 1;
    while(N){
        i64 n = N % P;
        vector<Modint> A(P-1);
        rep(i,n+1) A[discreteLog[comb(n, i).val()]] += 1;
        res = cycomb(res, A);
        N /= P;
    }
    Modint ans = 0;
    rep(i,P) if(i) ans += res[discreteLog[i]] * i;
    cout << ans.val() << endl;
}

int main(){
    ios::sync_with_stdio(false); cin.tie(nullptr);
    #ifdef NACHIA
    int T; cin >> T; for(int t=0; t<T; T!=++t?(cout<<'\n'),0:0)
    #endif
    testcase();
    return 0;
}
yukicoder

結果

ソースコード
