#1008545 (C++17) No.2883 K-powered Sum of Fibonacci

提出ソース

結果

問題	No.2883 K-powered Sum of Fibonacci
ユーザー	nonon
提出日時	2024-09-08 15:35:58
言語	C++17 (gcc 12.3.0 + boost 1.83.0)
結果	AC
実行時間	4 ms / 3,000 ms
コード長	15,668 bytes
コンパイル時間	3,003 ms
コンパイル使用メモリ	228,564 KB
実行使用メモリ	6,944 KB
最終ジャッジ日時	2024-09-08 15:36:03
合計ジャッジ時間	4,408 ms
ジャッジサーバーID （参考情報）	judge4 / judge3

このコードへのチャレンジ
（要ログイン）

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	2 ms 6,816 KB
testcase_01	AC	1 ms 6,820 KB
testcase_02	AC	2 ms 6,816 KB
testcase_03	AC	2 ms 6,940 KB
testcase_04	AC	2 ms 6,944 KB
testcase_05	AC	2 ms 6,940 KB
testcase_06	AC	2 ms 6,940 KB
testcase_07	AC	2 ms 6,940 KB
testcase_08	AC	2 ms 6,940 KB
testcase_09	AC	2 ms 6,940 KB
testcase_10	AC	2 ms 6,940 KB
testcase_11	AC	3 ms 6,940 KB
testcase_12	AC	3 ms 6,944 KB
testcase_13	AC	2 ms 6,944 KB
testcase_14	AC	3 ms 6,940 KB
testcase_15	AC	3 ms 6,944 KB
testcase_16	AC	2 ms 6,940 KB
testcase_17	AC	2 ms 6,944 KB
testcase_18	AC	2 ms 6,940 KB
testcase_19	AC	3 ms 6,944 KB
testcase_20	AC	4 ms 6,940 KB
testcase_21	AC	4 ms 6,940 KB
testcase_22	AC	4 ms 6,944 KB
testcase_23	AC	3 ms 6,944 KB
testcase_24	AC	4 ms 6,944 KB
testcase_25	AC	3 ms 6,940 KB
testcase_26	AC	4 ms 6,944 KB
testcase_27	AC	3 ms 6,944 KB
testcase_28	AC	4 ms 6,944 KB
testcase_29	AC	3 ms 6,944 KB
testcase_30	AC	2 ms 6,940 KB
testcase_31	AC	2 ms 6,944 KB
testcase_32	AC	1 ms 6,940 KB
testcase_33	AC	2 ms 6,940 KB
testcase_34	AC	2 ms 6,940 KB
testcase_35	AC	2 ms 6,944 KB
testcase_36	AC	2 ms 6,944 KB
testcase_37	AC	2 ms 6,940 KB
testcase_38	AC	2 ms 6,944 KB
testcase_39	AC	2 ms 6,940 KB
testcase_40	AC	2 ms 6,944 KB
testcase_41	AC	2 ms 6,940 KB
testcase_42	AC	4 ms 6,944 KB

権限があれば一括ダウンロードができます

ソースコード

raw source code

#include<bits/stdc++.h>
#include<atcoder/modint>
using namespace std;
using mint=atcoder::modint998244353;
template<typename mint>
struct Number_Theoretic_Transform
{
    static vector<mint>dw,dw_inv;
    static int log;
    static mint root;
    static void ntt(vector<mint>& f)
    {
        init();
        const int n=f.size();
        for(int m=n;m>>=1;)
        {
            mint w=1;
            for(int s=0,k=0;s<n;s+=(m<<1))
            {
                for(int i=s,j=s+m;i<s+m;i++,j++)
                {
                    mint x=f[i],y=f[j]*w;
                    f[i]=x+y,f[j]=x-y;
                }
                w*=dw[__builtin_ctz(++k)];
            }
        }
    }
    static void intt(vector<mint>& f, bool flag=true)
    {
        init();
        const int n=f.size();
        for(int m=1;m<n;m<<=1)
        {
            mint w=1;
            for(int s=0,k=0;s<n;s+=(m<<1))
            {
                for(int i=s,j=s+m;i<s+m;i++,j++)
                {
                    mint x=f[i],y=f[j];
                    f[i]=x+y,f[j]=(x-y)*w;
                }
                w*=dw_inv[__builtin_ctz(++k)];
            }
        }
        if(flag)
        {
            mint cef=mint(n).inv();
            for(int i=0;i<n;i++)f[i]*=cef;
        }
    }
private:
    Number_Theoretic_Transform()=default;
    static void init()
    {
        if(!dw.empty())return;
        long long mod=998244353;
        long long tmp=mod-1;
        log=1;
        while(tmp%2==0)
        {
            tmp>>=1;
            log++;
        }
        dw.resize(log);
        dw_inv.resize(log);
        for(int i=0;i<log;i++)
        {
            dw[i]=-root.pow((mod-1)>>(i+2));
            dw_inv[i]=dw[i].inv();
        }
    }
};
template<typename mint>
vector<mint>Number_Theoretic_Transform<mint>::dw=vector<mint>();
template<typename mint>
vector<mint>Number_Theoretic_Transform<mint>::dw_inv=vector<mint>();
template<typename mint>
int Number_Theoretic_Transform<mint>::log=0;
template<typename mint>
mint Number_Theoretic_Transform<mint>::root=mint(3);
template<typename mint>
struct Formal_Power_Series:vector<mint>
{
    using FPS=Formal_Power_Series;
    using vector<mint>::vector;
    using NTT=Number_Theoretic_Transform<mint>;
    void ntt(){NTT::ntt(*this);}
    void intt(bool flag=true){NTT::intt(*this,flag);}
    Formal_Power_Series(const vector<mint>&x, const vector<mint>&y)
    {
        const int n=x.size();
        auto rec=[&](const auto&rec, int l, int r)->FPS
        {
            if(l+1==r)return FPS{-x[l],1};
            int m=(l+r)/2;
            return rec(rec,l,m)*rec(rec,m,r);
        };
        FPS g=rec(rec,0,n);
        vector<mint> dg=g.diff().multipoint_evaluation(x);
        auto rec2=[&](const auto&rec2, int l, int r)->pair<FPS,FPS>
        {
            if(l+1==r)return make_pair(FPS{y[l]/dg[l]},FPS{-x[l],1});
            int m=(l+r)/2;
            auto[fl,gl]=rec2(rec2,l,m);
            auto[fr,gr]=rec2(rec2,m,r);
            return make_pair(fl*gr+fr*gl,gl*gr);
        };
        *this=rec2(rec2,0,n).first;
    }
    FPS &operator+=(const mint& r)
    {
        if(this->empty())this->resize(1);
        (*this)[0]+=r;
        return *this;
    }
    FPS &operator-=(const mint& r)
    {
        if(this->empty())this->resize(1);
        (*this)[0]-=r;
        return *this;
    }
    FPS &operator*=(const mint& r)
    {
        for(mint &x:*this)x*=r;
        return *this;
    }
    FPS &operator/=(const mint& r)
    {
        mint invr=r.inv();
        for(mint &x:*this)x*=invr;
        return *this;
    }
    FPS operator+(const mint& r)const{return FPS(*this)+=r;}
    FPS operator-(const mint& r)const{return FPS(*this)-=r;}
    FPS operator*(const mint& r)const{return FPS(*this)*=r;}
    FPS operator/(const mint& r)const{return FPS(*this)/=r;}
    FPS& operator+=(const FPS& f)
    {
        if(this->size()<f.size())this->resize(f.size());
        for(int i=0;i<(int)f.size();i++)(*this)[i]+=f[i];
        return *this;
    }
    FPS& operator-=(const FPS& f)
    {
        if(this->size()<f.size())this->resize(f.size());
        for(int i=0;i<(int)f.size();i++)(*this)[i]-=f[i];
        return *this;
    }
    FPS& operator*=(const FPS& f)
    {
        *this=convolution(*this,f);
        return *this;
    }
    FPS& operator/=(const FPS& f)
    {
        return *this*=f.inv();
    }
    FPS& operator%=(const FPS& f)
    {
        *this-=this->div(f)*f;
        this->shrink();
        return *this;
    }
    FPS operator+(const FPS& f)const{return FPS(*this)+=f;}
    FPS operator-(const FPS& f)const{return FPS(*this)-=f;}
    FPS operator*(const FPS& f)const{return FPS(*this)*=f;}
    FPS operator/(const FPS& f)const{return FPS(*this)/=f;}
    FPS operator%(const FPS& f)const{return FPS(*this)%=f;}
    FPS operator-()const
    {
        FPS res(this->size());
        for(int i=0;i<(int)this->size();i++)res[i]-=(*this)[i];
        return res;
    }
    FPS div(FPS f)const
    {
        if(this->size()<f.size())return FPS{};
        int n=this->size()-f.size()+1;
        return (rev().pre(n)*f.rev().inv(n)).pre(n).rev(n);
    }
    FPS pre(int deg)const
    {
        return FPS(begin(*this),begin(*this)+min((int)this->size(),deg));
    }
    FPS rev(int deg=-1)const
    {
        FPS res(*this);
        if(deg!=-1)res.resize(deg,0);
        reverse(begin(res),end(res));
        return res;
    }
    void shrink()
    {
        while(!this->empty()&&this->back()==0)this->pop_back();
    }
    FPS dot(FPS f)const
    {
        int n=min(this->size(),f.size());
        FPS res(n);
        for(int i=0;i<n;i++)res[i]=(*this)[i]*f[i];
        return res;
    }
    FPS operator<<(int deg)const
    {
        FPS res(*this);
        res.insert(res.begin(),deg,0);
        return res;
    }
    FPS& operator<<=(int deg)
    {
        return *this=*this<<(deg);
    }
    FPS operator>>(int deg)const
    {
        if((int)this->size()<=deg)return{};
        FPS res(*this);
        res.erase(res.begin(),res.begin()+deg);
        return res;
    }
    FPS& operator>>=(int deg)
    {
        return *this=*this>>(deg);
    }
    mint operator()(const mint& r)const
    {
        mint res=0,powr=1;
        for(auto x:*this)
        {
            res+=x*powr;
            powr*=r;
        }
        return res;
    }
    FPS diff()const
    {
        int n=this->size();
        FPS res(max(0,n-1));
        for(int i=1;i<n;i++)
        {
            res[i-1]=(*this)[i]*i;
        }
        return res;
    }
    FPS integral()const
    {
        int n=this->size();
        FPS res(n+1);
        res[0]=0;
        for(int i=0;i<n;i++)
        {
            res[i+1]=(*this)[i]/(i+1);
        }
        return res;
    }
    FPS inv(int deg=-1)const
    {
        assert(((*this)[0])!=(0));
        int n=this->size();
        if(deg==-1)deg=n;
        FPS res(deg);
        res[0]={(*this)[0].inv()};
        for(int d=1;d<deg;d<<=1)
        {
            FPS f(d<<1),g(d<<1);
            for(int j=0;j<min(n,2*d);j++)f[j]=(*this)[j];
            for(int j=0;j<d;j++)g[j]=res[j];
            f.ntt();
            g.ntt();
            f=f.dot(g);
            f.intt();
            for(int j=0;j<d;j++)f[j]=0;
            f.ntt();
            f=f.dot(g);
            f.intt();
            for(int j=d;j<min(2*d,deg);j++)res[j]=-f[j];
        }
        return res;
    }
    FPS exp(int deg=-1)const
    {
        assert((*this)[0]==0);
        if(deg==-1)deg=this->size();
        vector<mint>inv;
        inv.reserve(deg+1);
        inv.push_back(mint::raw(0));
        inv.push_back(mint::raw(1));
        auto inplace_integral=[&](FPS& f)->void
        {
            int n=f.size();
            long long mod=mint::mod();
            while(inv.size()<=f.size())
            {
                int i=inv.size();
                inv.push_back((-inv[mod%i])*(mod/i));
            }
            f.insert(begin(f),mint::raw(0));
            for(int i=1;i<=n;i++)f[i]*=inv[i];
        };
        auto inplace_diff=[](FPS& f)->void
        {
            if(f.empty())return;
            f.erase(begin(f));
            mint cef=1;
            for(int i=0;i<(int)f.size();i++)
            {
                f[i]*=cef;
                cef++;
            }
        };
        FPS b={1,1<this->size()?(*this)[1]:0};
        FPS c={1},z1,z2={1,1};
        for(int m=2;m<deg;m<<=1)
        {
            FPS y=b;
            y.resize(2*m);
            y.ntt();
            z1=z2;
            FPS z(m);
            z=y.dot(z1);
            z.intt();
            fill(begin(z),begin(z)+m/2,mint::raw(0));
            z.ntt();
            z=z.dot(-z1);
            z.intt();
            c.insert(end(c),begin(z)+m/2,end(z));
            z2=c;
            z2.resize(2*m);
            z2.ntt();
            FPS x(begin(*this),begin(*this)+min(int(this->size()),m));
            inplace_diff(x);
            x.push_back(mint::raw(0));
            x.ntt();
            x=x.dot(y);
            x.intt();
            x-=b.diff();
            x.resize(2*m);
            for(int i=0;i<m-1;i++)x[m+i]=x[i],x[i]=0;
            x.ntt();
            x=x.dot(z2);
            x.intt();
            x.pop_back();
            inplace_integral(x);
            for(int i=m;i<min(int(this->size()),2*m);i++)x[i]+=(*this)[i];
            fill(begin(x),begin(x)+m,mint::raw(0));
            x.ntt();
            x=x.dot(y);
            x.intt();
            b.insert(end(b),begin(x)+m,end(x));
        }
        return FPS{begin(b),begin(b)+deg};
    }
    FPS log(int deg=-1)const
    {
        assert((*this)[0]==1);
        int n=this->size();
        if(deg==-1)deg=n;
        return (this->diff()*this->inv()).pre(deg-1).integral();
    }
    FPS pow(long long k, int deg=-1)const
    {
        if(deg==-1)deg=this->size();
        if(k==0)
        {
            FPS res(deg);
            res[0]=mint::raw(1);
            return res;
        }
        FPS res=*this;
        int cnt0=0;
        while(cnt0<res.size()&&res[cnt0]==0)cnt0++;
        if (cnt0>(deg-1)/k)
        {
            FPS res(deg);
            return res;
        }
        res=res>>cnt0;
        deg-=cnt0*k;
        res=((res/res[0]).log(deg)*k).exp(deg)*res[0].pow(k);
        res=res<<(cnt0*k);
        return res.pre(deg);
    }
    FPS sqrt(int deg=-1)const
    {
        auto sqrt_mod=[](mint r)->mint
        {
            const int mod=mint::mod();
            if(r==0||r==1)return r;
            if(r.pow((mod-1)>>1)!=1)return -1;
            mint b=1;
            while(b.pow((mod-1)>>1)==1)b++;
            int m=mod-1,e=0;
            while((m&1)==0)m>>=1,e++;
            mint x=r.pow((m-1)>>1);
            mint y=r*x*x;
            x*=r;
            mint z=b.pow(m);
            while(y!=1)
            {
                int j=0;
                mint t=y;
                while(t!=1)j++,t*=t;
                z=z.pow(1LL<<(e-j-1));
                x*=z;
                z*=z;
                y*=z;
                e=j;
            }
            return x;
        };
        if(deg==-1)deg=this->size();
        int cnt0=0;
        while(cnt0<this->size()&&(*this)[cnt0]==0)cnt0++;
        if(cnt0)
        {
            if(cnt0==(int)this->size())return FPS(deg);
            if(cnt0&1)return{};
            if(2*deg<=cnt0)return FPS(deg);
            FPS res=(*this>>cnt0).sqrt(deg-cnt0/2);
            if(res.empty())return{};
            res=res<<(cnt0/2);
            res.resize(deg,0);
            return res;
        }
        mint sqr=sqrt_mod((*this)[0]);
        if(sqr*sqr!=(*this)[0])return{};
        FPS res={sqr};
        mint inv2=mint(2).inv();
        for(int i=1;i<deg;i<<=1)res=(res+this->pre(i<<1)*res.inv(i<<1))*inv2;
        return res.pre(deg);
    }
    FPS circular_mod(int m)const
    {
        FPS res(m);
        for(int i=0;i<(int)this->size();i++)res[i%m]+=(*this)[i];
        return res;
    }
    FPS taylor_shift(mint c)const
    {
        int n=this->size();
        FPS fact(n),fact_inv(n);
        { // calc fact and fact inv
            fact[0]=1;
            for(int i=1;i<n;i++)fact[i]=i*fact[i-1];
            fact_inv[n-1]=fact[n-1].inv();
            for(int i=n-1;i>=1;i--)fact_inv[i-1]=i*fact_inv[i];
        }
        FPS res(*this);
        res=res.dot(fact);
        res=res.rev();
        FPS bs(n,mint::raw(1));
        for(int i=1;i<n;i++)bs[i]=bs[i-1]*c*fact_inv[i]*fact[i-1];
        res=(res*bs).pre(n);
        res=res.rev();
        res=res.dot(fact_inv);
        return res;
    }
    vector<mint>multipoint_evaluation(const vector<mint>&x)const
    {
        if(x.empty())return{};
        int m=x.size(),n=1;
        if(this->size()==0){return vector<mint>(m,0);}
        if(this->size()==1){return vector<mint>(m,(*this)[0]);}
        while(m>n)n<<=1;
        vector<FPS>f(n<<1,FPS({mint(1)}));
        for(int i=0;i<m;i++)f[i+n]=FPS({-x[i],mint(1)});
        for(int i=n-1;i>0;i--)f[i]=f[i<<1]*f[(i<<1)|1];
        f[1]=(*this)%f[1];
        for(int i=2;i<n+m;i++)f[i]=f[i>>1]%f[i];
        vector<mint>res(m);
        for(int i=0;i<m;i++)res[i]=(f[i+n].empty()?mint(0):f[i+n][0]);
        return res;
    }
private:
    FPS convolution(FPS f, FPS g)
    {
        int n=f.size(),m=g.size();
        if(n==0||m==0)return {};
        int log=1;
        while((1<<log)<n+m-1)log++;
        int sz=1<<log;
        f.resize(sz);
        g.resize(sz);
        f.ntt();
        g.ntt();
        mint inv=mint(sz).inv();
        for(int i=0;i<sz;i++)f[i]*=g[i]*inv;
        f.intt(0);
        f.resize(n+m-1);
        return f;
    }
};
using FPS=Formal_Power_Series<mint>;
vector<mint>Berlekamp_Massey(const vector<mint>&a)
{
    int n=a.size();
    vector<mint>b,c;
    b.reserve(n+1);b.push_back(1);
    c.reserve(n+1);c.push_back(1);
    mint y=1;
    for(int d=1;d<=n;d++)
    {
        int k=b.size(),l=c.size();
        mint x=0;
        for(int i=0;i<l;i++)x+=c[i]*a[d-l+i];
        b.push_back(0);
        k++;
        if(x==0)continue;
        mint buf=x/y;
        if(l<k)
        {
            vector<mint>tmp=c;
            c.insert(c.begin(),k-l,0);
            for(int i=0;i<k;i++)c[k-i-1]-=buf*b[k-i-1];
            b=tmp;
            y=x;
        }
        else
        {
            for(int i=0;i<k;i++)c[l-i-1]-=buf*b[k-i-1];
        }
    }
    reverse(c.begin(),c.end());
    for(mint& x:c)x=-x;
    return c;
}
mint Bostan_Mori(FPS p, FPS q, long long k)
{
    mint res=0;
    if(p.size()>=q.size())
    {
        FPS r=p.div(q);
        p-=q*r;
        p.shrink();
        if(k<(int)r.size())res+=r[k];
    }
    if(p.empty())return res;
    p.resize(q.size()-1);
    auto sub=[&](const FPS& f, bool odd=0)
    {
        FPS g((f.size()+!odd)/2);
        for(int i=odd;i<(int)f.size();i+=2)g[i>>1]=f[i];
        return g;
    };
    while(k)
    {
        FPS q2=q;
        for(int i=1;i<(int)q2.size();i+=2)q2[i]=-q2[i];
        p=sub(p*q2,k&1);
        q=sub(q*q2);
        k>>=1;
    }
    return res+p[0];
}
mint linear_recurrence(FPS a, FPS c, long long k)
{
    assert(a.size()==c.size());
    c=FPS{1}-(c<<1);
    return Bostan_Mori((a*c).pre(a.size()),c,k);
}
mint BMBM(vector<mint>&x, long long k)
{
    vector<mint>tmp=Berlekamp_Massey(x);
    int n=tmp.size()-1;
    FPS a(n),c(n);
    for(int i=0;i<n;i++)
    {
        a[i]=x[i];
        c[i]=tmp[i+1];
    }
    return linear_recurrence(a,c,k);
}
int main()
{
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    long N,K;
    cin>>N>>K;
    int M=4*K;
    vector<mint>F(M);
    F[0]=F[1]=1;
    vector<mint>A(M);
    A[0]=1,A[1]=2;
    for(int i=2;i<M;i++)
    {
        F[i]=F[i-1]+F[i-2];
        A[i]=A[i-1]+F[i].pow(K);
    }
    cout<<BMBM(A,N-1).val()<<endl;
}

yukicoder

結果

テストケース

ソースコード