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

提出ソース
結果

問題	No.1145 Sums of Powers
ユーザー	eQe
提出日時	2023-08-14 18:44:21
言語	C++23 (gcc 12.3.0 + boost 1.83.0)
結果	AC
実行時間	299 ms / 2,000 ms
コード長	15,997 bytes
コンパイル時間	7,542 ms
コンパイル使用メモリ	335,292 KB
実行使用メモリ	18,644 KB
最終ジャッジ日時	2024-05-02 06:47:55
合計ジャッジ時間	8,772 ms
ジャッジサーバーID （参考情報）	judge1 / judge3
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 5,248 KB
testcase_01	AC	2 ms 5,376 KB
testcase_02	AC	4 ms 5,376 KB
testcase_03	AC	299 ms 17,824 KB
testcase_04	AC	299 ms 18,644 KB
testcase_05	AC	295 ms 17,932 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
#include<bits/stdc++.h>
#include<atcoder/all>
#define eb emplace_back
#define all(a) (a).begin(),(a).end()
#define RD(T,...) T __VA_ARGS__;li(__VA_ARGS__)
#define LL(...) RD(ll,__VA_ARGS__)
#define VL(n,...) vl __VA_ARGS__;setsize({n},__VA_ARGS__);li(__VA_ARGS__)
#define JO(a,b) a##b
#define jo(a,b) JO(a,b)
#define uq_symbol(a) jo(a,__LINE__)
#define FE(v,e,...) for(au&&__VA_OPT__([)e __VA_OPT__(,__VA_ARGS__]):v)
#define FO(n) for(ll uq_symbol(i)=n;uq_symbol(i)-->0;)
#define FOR(i,...) for(au[i,i##O,i##E]=rng(0,__VA_ARGS__);i<i##O;i+=i##E)
#define fe(v,...) FE(v,__VA_ARGS__)
#define fo(i,...) FO##__VA_OPT__(R)(i __VA_OPT__(,__VA_ARGS__))
#define of(i,...) for(au[i,i##O,i##E]=rng(1,__VA_ARGS__);i>=i##O;i-=i##E)
#define I template
#define J class
#define au auto
#define rr return
using std::cin,std::cout,std::begin,std::end,std::rbegin;
using std::swap,std::move,std::abs,std::prev,std::next;
using std::tuple,std::array,std::bitset,std::minmax,std::get;
using vo=void;using bo=bool;
vo solve();
using is=std::istream;using os=std::ostream;
using dd=long double;using ll=long long;using ull=unsigned long long;
using lll=__int128_t;using ulll=__uint128_t;using str=std::string;
using namespace atcoder;using ml=modint;au&operator<<(os&o,const ml&x){rr o<<x.val();}
os&operator<<(os&o,const ulll&x){rr(x<10?o:o<<x/10)<<ll(x%10);}
os&operator<<(os&o,const lll&x){rr o<<str(x<0,'-')<<ulll(x>0?x:-x);}
constexpr ll oo=3e18;
constexpr ll dx[]{-1,0,1,0,-1,1,1,-1},dy[]{0,-1,0,1,-1,-1,1,1};
constexpr char sp=32;
constexpr char nl=10;
au rng(bo s,ll a,ll b=oo,ll c=1){if(b==oo)b=a,(s?b:a)=0;rr tuple{a-s,b,c};}

ll pwm1(ll x){rr 1-2*(x&1);}
I<J T>T sq(T a){rr a*a;}
I<J T>ll len(const T&a){rr a.size();}
I<J...A>au max(const A&...a){rr max(std::initializer_list<std::common_type_t<A...>>{a...});}
I<J...A>au min(const A&...a){rr min(std::initializer_list<std::common_type_t<A...>>{a...});}

struct edg{
  ll t,w;
  edg(){}
  edg(ll t,ll w=1):t(t),w(w){}
  friend os&operator<<(os&o,const edg&e){rr o<<e.t<<sp<<e.w;}
};

I<J A,J B=A>struct cp{
  A a={};B b={};
  cp(){}
  cp(A a,B b):a(a),b(b){}
  cp(std::pair<A,B>p):a(p.first),b(p.second){}
  bo operator==(const cp&c)const{rr a==c.a&&b==c.b;}
  bo operator<(const cp&c)const{rr a!=c.a?a<c.a:b<c.b;}
  bo operator>(const cp&c)const{rr a!=c.a?a>c.a:b>c.b;}

  friend is&operator>>(is&i,cp&c){rr i>>c.a>>c.b;}
  friend os&operator<<(os&o,const cp&c){rr o<<c.a<<sp<<c.b;}
};using cl=cp<ll>;
I<J A,J B=A,J C=A>struct tr{
  A a={};B b={};C c={};
  tr(){}
  tr(A a,B b,C c):a(a),b(b),c(c){}
  bo operator==(const tr&t)const{rr a==t.a&&b==t.b&&c==t.c;}

  au operator<(const tr&t)const{rr a!=t.a?a<t.a:b!=t.b?b<t.b:c<t.c;}
  friend is&operator>>(is&i,tr&t){rr i>>t.a>>t.b>>t.c;}
  friend os&operator<<(os&o,const tr&t){rr o<<t.a<<sp<<t.b<<sp<<t.c;}
};using tl=tr<ll>;

I<J T>T&rv(T&a){reverse(all(a));rr a;}
I<J T>au pot(T&a){au r=a.top();a.pop();rr r;}

I<J T>struct qmin:std::priority_queue<T,std::vector<T>,std::greater<>>{
  qmin(const std::vector<T>&a={}){fe(a,e)this->emplace(e);}
  qmin(const std::initializer_list<T>&a){fe(a,e)this->emplace(e);}
  friend os&operator<<(os&o,qmin q){while(len(q))o<<pot(q)<<str(len(q)>0,sp);rr o;}
};

I<J V>struct ve;
I<J V>constexpr bo isv=0;
I<J V>constexpr bo isv<ve<V>> =1;
I<J V>constexpr bo isv<std::vector<V>> =1;
I<J V>au rawv(V){if constexpr(isv<V>)rr rawv(V(1)[0]);else rr V();}

I<J V>struct ve:std::vector<V>{
  using std::vector<V>::vector;
  using T=decltype(rawv(V()));
  I<J U>ve(const std::vector<U>&v={}){fe(v,e)this->eb(e);}

  ve&operator+=(const ve&u){au&v=*this;fo(i,len(v))v[i]+=u[i];rr v;}
  ve&operator-=(const ve&u){au&v=*this;fo(i,len(v))v[i]-=u[i];rr v;}
  ve&operator^=(const ve&u){fe(u,e)this->eb(e);rr*this;}
  ve&operator+=(const T&x){au&v=*this;fe(v,e)e+=x;rr v;}
  ve&operator-=(const T&x){au&v=*this;fe(v,e)e-=x;rr v;}
  ve&operator*=(const T&x){au&v=*this;fe(v,e)e*=x;rr v;}
  I<size_t n>au&operator+=(const bitset<n>&a){fo(i,n)(*this)[i]+=a[i];rr*this;}
  I<size_t n>au&operator-=(const bitset<n>&a){fo(i,n)(*this)[i]-=a[i];rr*this;}

  ve operator+(const ve&u)const{rr ve(*this)+=u;}
  ve operator-(const ve&u)const{rr ve(*this)-=u;}
  ve operator^(const ve&u)const{rr ve(*this)^=u;}
  ve operator+(const T&x)const{rr ve(*this)+=x;}
  ve operator-(const T&x)const{rr ve(*this)-=x;}
  ve operator*(const T&x)const{rr ve(*this)*=x;}
  ve&operator++(){rr*this+=1;}
  ve&operator--(){rr*this-=1;}
  ve operator-()const{rr ve(*this)*=-1;}

  au lower_bound(const V&x)const{rr std::lower_bound(all(*this),x);}
  au upper_bound(const V&x)const{rr std::upper_bound(all(*this),x);}

  I<J F>au scan(F f)const{cp<T,bo>r;fe(*this,e)if constexpr(!isv<V>)r.b?f(r.a,e),r:r={e,1};else if(au s=e.scan(f);s.b)r.b?f(r.a,s.a),r:r=s;rr r;}
  T max()const{rr scan([](T&a,const T&b){a<b?a=b:0;}).a;}
  T min()const{rr scan([](T&a,const T&b){a>b?a=b:0;;}).a;}

  vo df(){this->erase(this->begin());}

};
I<J T=ll,size_t n,size_t i=0>au vec(const ll(&s)[n],T x={}){
  if constexpr(n==i+1)rr ve<T>(s[i],x);
  else{au X=vec<T,n,i+1>(s,x);rr ve<decltype(X)>(s[i],X);}
}
I<ll n,J...A>vo setsize(const ll(&l)[n],A&...a){((a=vec(l,rawv(a))),...);}
I<J T>using vve=ve<ve<T>>;using vl=ve<ll>;using vvl=vve<ll>;
using gl=ve<ve<ll>>;using ge=ve<ve<edg>>;

struct io{io(){cin.tie(0)->sync_with_stdio(0);
  cout<<std::fixed<<std::setprecision(15);std::cerr<<nl;}}io;
I<size_t n>os&operator<<(os&o,const bitset<n>&b){fo(i,n)o<<b[i];rr o;}
I<J...A>os&operator<<(os&o,const tuple<A...>&t){
  apply([&](const au&...a){ll i=0;(((o<<a<<str(++i!=sizeof...(a),sp))),...);},t);rr o;}
I<J V>os&operator<<(os&o,const std::vector<V>&v){fe(v,e)o<<e<<str(&e!=&v.back(),isv<V>?nl:sp);rr o;}
I<char c=sp,J...A>vo pp(const A&...a){ll i=0;((cout<<a<<str(++i!=sizeof...(a),c)),...);cout<<nl;}
I<J V>is&operator>>(is&i,std::vector<V>&v){fe(v,e)i>>e;rr i;}
I<J...A>vo li(A&...a){(cin>>...>>a);}

ll inv(ll x,ll m){ll a=(x%m+m)%m,b=m,u=1,v=0;while(b)u-=a/b*v,swap(u,v),a-=a/b*b,swap(a,b);rr(u%m+m)%m;}
I<J F>struct rec:F{rec(F&&f):F(std::forward<F>(f)){}I<J...A>decltype(au)operator()(A&&...a)const{rr F::operator()(*this,std::forward<A>(a)...);}};

//https://ei1333.github.io/luzhiled/snippets/math/fast-fourier-transform.html
namespace fft{
using real=double;
struct C{
  real x,y;
  C():x(0),y(0){}
  C(real x,real y):x(x),y(y){}
  inline C operator+(const C &c)const{return C(x+c.x,y+c.y);}
  inline C operator-(const C &c)const{return C(x-c.x,y-c.y);}
  inline C operator*(const C &c)const{return C(x*c.x-y*c.y,x*c.y+y*c.x);}
  inline C conj()const{return C(x,-y);}
};

const real PI=acosl(-1);
ll base=1;
std::vector<C>rts={{0,0},{1,0}};
std::vector<int>rev={0,1};

vo ensure_base(int nbase){
  if(nbase<=base)return;
  rev.resize(1<<nbase);
  rts.resize(1<<nbase);
  fo(i,1<<nbase)rev[i]=(rev[i>>1]>>1)+((i&1)<<(nbase-1));
  while(base<nbase){
    real angle=PI*2.0/(1<<(base+1));
    fo(i,1<<(base-1),1<<base){
      rts[i<<1]=rts[i];
      real angle_i=angle*(2*i+1-(1<<base));
      rts[(i<<1)+1]=C(std::cos(angle_i),std::sin(angle_i));
    }
    ++base;
  }
}

vo fft(std::vector<C>&a,int n){
  assert((n&(n-1))==0);
  int zeros=__builtin_ctz(n);
  ensure_base(zeros);
  int shift=base-zeros;
  fo(i,n)if(i<(rev[i]>>shift))swap(a[i],a[rev[i]>>shift]);

  for(int k=1;k<n;k<<=1){
    for(int i=0;i<n;i+=2*k){
      for(int j=0;j<k;j++){
        C z=a[i+j+k]*rts[j+k];
        a[i+j+k]=a[i+j]-z;
        a[i+j]=a[i+j]+z;
      }
    }
  }
}
};

//https://ei1333.github.io/luzhiled/snippets/math/arbitrary-mod-convolution.html
template<class T>struct arbitrary_mod_convolution{
  using real=fft::real;
  using C=fft::C;
  arbitrary_mod_convolution()=default;

  std::vector<T>multiply(const std::vector<T>&a,const std::vector<T>&b,int need=-1){
    if(need==-1)need=len(a)+len(b)-1;
    int nbase=0;
    while((1<<nbase)<need)nbase++;
    fft::ensure_base(nbase);
    int sz=1<<nbase;
    std::vector<C>fa(sz);
    fo(i,len(a))fa[i]=C(a[i].val()&((1<<15)-1),a[i].val()>>15);

    fft::fft(fa,sz);
    std::vector<C>fb(sz);
    if(a==b){
      fb=fa;
    }else{
      fo(i,len(b))fb[i]=C(b[i].val()&((1<<15)-1),b[i].val()>>15);
      fft::fft(fb,sz);
    }
    real ratio=0.25/sz;
    C r2(0,-1),r3(ratio,0),r4(0,-ratio),r5(0,1);
    for(int i=0;i<=(sz>>1);i++){
      int j=(sz-i)&(sz-1);
      C a1=(fa[i]+fa[j].conj());
      C a2=(fa[i]-fa[j].conj())*r2;
      C b1=(fb[i]+fb[j].conj())*r3;
      C b2=(fb[i]-fb[j].conj())*r4;
      if(i!=j){
        C c1=(fa[j]+fa[i].conj());
        C c2=(fa[j]-fa[i].conj())*r2;
        C d1=(fb[j]+fb[i].conj())*r3;
        C d2=(fb[j]-fb[i].conj())*r4;
        fa[i]=c1*d1+c2*d2*r5;
        fb[i]=c1*d2+c2*d1;
      }
      fa[j]=a1*b1+a2*b2*r5;
      fb[j]=a1*b2+a2*b1;
    }
    fft::fft(fa,sz);
    fft::fft(fb,sz);
    std::vector<T>ret(need);
    fo(i,need){
      int64_t aa=llround(fa[i].x);
      int64_t bb=llround(fb[i].x);
      int64_t cc=llround(fa[i].y);
      aa=T(aa).val(),bb=T(bb).val(),cc=T(cc).val();
      ret[i]=aa+(bb<<15)+(cc<<30);
    }
    return ret;
  }
};
using amc=arbitrary_mod_convolution<ml>;

//https://ei1333.github.io/luzhiled/snippets/math/formal-power-series.html
template<class T>struct formal_power_series:ve<T>{
  using ve<T>::ve;
  using P=formal_power_series;

  using MULT=std::function<P(P,P)>;

  static MULT&get_mult(){static MULT mult=nullptr;return mult;}
  static vo set_fft(MULT f){get_mult()=f;}

  template<class U>formal_power_series(const ve<U>&a){
    ll n=len(a);
    this->resize(n);
    fo(i,n)(*this)[i]=a[i];
  }

  bo operator<(const P&f)const{return len(*this)<len(f);}
  bo operator>(const P&f)const{return len(*this)>len(f);}

  P&operator+=(const P&f){
    if(len(f)>len(*this))this->resize(len(f));
    fo(i,len(f))(*this)[i]+=f[i];
    return*this;
  }
  P&operator-=(const P&f){
    if(len(f)>len(*this))this->resize(len(f));
    fo(i,len(f))(*this)[i]-=f[i];
    return*this;
  }

  P&operator+=(const T&t){if(!len(*this))this->resize(1);(*this)[0]+=t;return*this;}
  P&operator-=(const T&t){if(!len(*this))this->resize(1);(*this)[0]-=t;return*this;}
  P&operator*=(const T&t){fo(i,len(*this))(*this)[i]*=t;return*this;}

  P&operator*=(const P&f){
    if(!len(*this)||!len(f))return*this=P();
    assert(get_mult());
    return*this=get_mult()(*this,f);
  }
  P&operator%=(const P&f){*this-=*this/f*f;shrink();return*this;}

  P&operator/=(const P&f){
    if(len(*this)<len(f))return*this=P();
    ll n=len(*this)-len(f)+1;
    return*this=(rev().pre(n)*f.rev().inv(n)).pre(n).rev(n);
  }

  P operator+(const P&f)const{return P(*this)+=f;}
  P operator-(const P&f)const{return P(*this)-=f;}
  P operator*(const P&f)const{return P(*this)*=f;}
  P operator/(const P&f)const{return P(*this)/=f;}
  P operator%(const P&f)const{return P(*this)%=f;}
  template<class U=T>P operator+(const U&t)const{return P(*this)+=t;}
  template<class U=T>P operator-(const U&t)const{return P(*this)-=t;}
  template<class U=T>P operator*(const U&t)const{return P(*this)*=t;}
  P operator-()const{P r=*this;fe(r,x)x=-x;return r;}
  P operator>>(ll sz)const{if(len(*this)<=sz)return P();P r(*this);r.erase(begin(r),begin(r)+sz);return r;}
  P operator<<(ll sz)const{P r(*this);r.insert(begin(r),sz,T{});return r;}
  vo shrink(){while(len(*this)&&this->back()==T{})this->pop_back();}
  P pre(ll deg)const{return P(this->begin(),this->begin()+min(len(*this),deg));}
  P rev(ll deg=-1)const{P r(*this);if(deg!=-1)r.resize(deg,T{});rv(r);return r;}
  T operator()(T x)const{T r=0,w=1;fe(*this,v)r+=w*v,w*=x;return r;}

  P diff()const{
    ll n=len(*this);
    P r(max(n-1,0));
    fo(i,1,n)r[i-1]=(*this)[i]*T{i};
    return r;
  }
  P integral()const{
    ll n=len(*this);
    P r(n+1);
    r[0]=T{};
    fo(i,n)r[i+1]=(*this)[i]/T{i+1};
    return r;
  }

  P inv(ll deg=-1)const{
    assert((*this)[0]!=T{});
    ll n=len(*this);
    if(deg==-1)deg=n;
    P r{T{1}/(*this)[0]};

    for(ll i=1;i<deg;i<<=1)r=(r+r-sq(r)*pre(i<<1)).pre(i<<1);
    return r.pre(deg);
  }

  P log(ll deg=-1)const{
    assert((*this)[0]==T{1});
    ll n=len(*this);
    if(deg==-1)deg=n;
    return(this->diff()*this->inv(deg)).pre(deg-1).integral();
  }

  P exp(ll deg=-1)const{
    assert((*this)[0]==T{});
    ll n=len(*this);
    if(deg==-1)deg=n;
    P r{1};
    for(ll i=1;i<deg;i<<=1)r=(r*(pre(i<<1)+T{1}-r.log(i<<1))).pre(i<<1);
    return r.pre(deg);
  }

  P sqr(ll deg=-1)const{
    ll n=len(*this);
    if(deg==-1)deg=n;

    if((*this)[0]==T{}){
      fo(i,1,n){
        if((*this)[i]!=T{}){
          if(i&1)return{};
          if(deg-i/2<=0)break;
          P r=(*this>>i).sqr(deg-i/2);
          if(mu(r))return{};
          r=r<<(i/2);
          if(len(r)<deg)r.resize(deg,T{});
          return r;
        }
      }
      return P(deg,T{});
    }

    ll s=msqr((*this)[0].val(),ml::mod());
    if(sq(s)!=(*this)[0])return{};

    P r{s};
    T iv2=T{1}/T{2};
    for(ll i=1;i<deg;i<<=1)r=(r+pre(i<<1)*r.inv(i<<1))*iv2;
    return r.pre(deg);
  }

  P pow(lll k,ll deg=-1)const{
    ll n=len(*this);
    if(deg==-1)deg=n;
    if(k==0){P r(deg);r[0]=1;return r;}

    fo(i,n){
      if((*this)[i]!=T{}){
        T rv=T{1}/(*this)[i];
        P r=(((*this*rv)>>i).log(deg)*k).exp(deg)*((*this)[i].pow(k));
        if(i*k>deg)return P(deg,T{});
        r=(r<<(i*k)).pre(deg);
        if(len(r)<deg)r.resize(deg,T{});
        return r;
      }
    }
    return*this;
  }

};
template<class T>constexpr bo isv<formal_power_series<T>> =1;

//https://nyaannyaan.github.io/library/fps/kitamasa.hpp.html
template<class T>T kitamasa(ll n,formal_power_series<T>p,formal_power_series<T>q){
  q.shrink();
  T ret=0;
  if(len(p)>=len(q)){
    formal_power_series<T>r=p/q;
    p-=r*q;
    p.shrink();
    if(n<len(r))ret+=r[n];
  }
  if(!len(p))return ret;

  p.resize(len(q)-1);
  while(n){
    formal_power_series<T>q2=q;
    fo(i,1,len(q2),2)q2[i]=-q2[i];
    formal_power_series<T>s=p*q2;
    formal_power_series<T>t=q*q2;
    if(n&1){
      fo(i,1,len(s),2)p[i>>1]=s[i];
      fo(i,0,len(t),2)q[i>>1]=t[i];
    }else{
      fo(i,0,len(s),2)p[i>>1]=s[i];
      fo(i,0,len(t),2)q[i>>1]=t[i];
    }
    n>>=1;
  }
  return ret+p[0];
}

template<class T>T kitamasa(ll n,const ve<T>&a,const ve<T>&c){
  assert(len(c)==len(a));
  ll k=len(c);
  formal_power_series<T>q=formal_power_series<T>{1}^-fps(c);
  formal_power_series<T>p=(q*fps(a)).pre(k);
  return kitamasa(n,p,q);
}

template<class T>formal_power_series<T>prod(const ve<formal_power_series<T>>&fs){
  qmin<formal_power_series<T>>q;
  fe(fs,f)q.emplace(f);
  while(len(q)>1){
    formal_power_series<T>f=q.top();q.pop();
    formal_power_series<T>g=q.top();q.pop();
    q.emplace(f*g);
  }
  return q.top();
}

template<class T>struct twelvefold{
  ve<T>fa,rfa,bs;
  vve<T>m;vve<int>u;
  twelvefold(ll n):fa(n+1,1),rfa(n+1,1),bs(n+1,1){
    fo(i,1,n+1)fa[i]=fa[i-1]*i;
    rfa[n]=fa[n].inv();
    of(i,n)rfa[i]=rfa[i+1]*(i+1);
    fo(i,1,n+1)bs[i]=bs[i-1]+pwm1(i)*rfa[i];
    if(n<=5000)m.resize(n+1,ve<int>(n+1)),u.resize(n+1,ve<int>(n+1));
  }
  T operator()(ll n,ll k){return c(n,k);}
  T c(ll n,ll k){return n<0?pwm1(k)*c(-n+k-1,k):k<0||n<k?0:fa[n]*rfa[k]*rfa[n-k];}
  T p(ll n,ll k){return c(n,k)*fa[k];}
  T h(ll n,ll r){return c(n-1+r,r);}
  T s(ll n,ll k){T r=0;fo(i,k+1)r+=pwm1(k-i)*c(k,i)*T(i).pow(n);return r*rfa[k];}
  T b(ll n,ll k){T r=0;fo(i,k+1)r+=T(i).pow(n)*rfa[i]*bs[k-i];return r;}
  T par(ll n,ll k){
    if(n==0&&k==0)return 1;
    if(k==0)return 0;
    if(u[n][k])return m[n][k];
    if(n-k>=0)return u[n][k]=1,m[n][k]=par(n,k-1)+par(n-k,k);
    return u[n][k]=1,m[n][k]=par(n,k-1);
  }
};using twf=twelvefold<ml>;

int main(){ll T=1;fo(T)solve();}
vo solve(){
  using ml=modint998244353;
  using fps=formal_power_series<ml>;
  amc fft;
  auto mul=[&](const fps&a,const fps&b){
    //auto r=fft.multiply(a,b);
    auto r=convolution(a,b);
    return fps(r);
  };
  fps::set_fft(mul);

  LL(N,M);
  VL(N,a);

  ve<fps>fs(N);
  fo(i,N)fs[i]=fps{1,-a[i]};
  fps f=prod(fs);
  f=-f.log(M+1);

  vl res(M+1);
  fo(i,1,M+1)res[i]=(f[i]*i).val();
  res.df();
  pp(res);
}
yukicoder

結果

テストケース

ソースコード
