math/kth_root.hpp

View this file on GitHub
Last update: 2021-01-30 11:24:14+09:00
Include: #include "math/kth_root.hpp"

Depends on

Code

#pragma once
#include"FPS_mint.hpp"
#include"mod_int_dynamic.hpp"

template<typename mint>
int kth_root(int n,int k){
    if(k==0){
        if(n==1)return 0;
        else return -1;
    }
    fps<mint>f(k+1,0);
    f[k]=1;
    f[0]=-n;
    random_device rnd;
    for(int times=0;times<10;++times){
        if(f.size()<=2){
            f.resize(k+1);
            f[k]=1;
            f[0]=-n;
        }
        fps<mint>g(k,0),h={1};
        for(int i=0;i<k;++i)g[i]=rnd();
        int t=(mint::get_mod()-1)/2;
        while(t){
            if(t%2)h*=g,h%=f,h.shrink();
            g*=g;
            g%=f;
            g.shrink();
            t/=2;
        }
        f=f.gcd(h+1).shrink();
        if(f.size()==2)return (f[0]/f[1]*(-1)).value();
    }
    return -1;
}

#line 2 "math/FPS_base.hpp"
#include<vector>
#include<tuple>
#include<iostream>
#include<cmath>
#include<type_traits>
#include<cassert>

/**
 * @brief 形式的冪級数(BASE)
 */

template<typename T,typename F>
struct FPS_BASE:std::vector<T>{
    using std::vector<T>::vector;
    using P=FPS_BASE<T,F>;
    F fft;
    FPS_BASE(){}
    inline P operator +(T x)const noexcept{return P(*this)+=x;}
    inline P operator -(T x)const noexcept{return P(*this)-=x;}
    inline P operator *(T x)const noexcept{return P(*this)*=x;}
    inline P operator /(T x)const noexcept{return P(*this)/=x;}
    inline P operator <<(int x)noexcept{return P(*this)<<=x;}
    inline P operator >>(int x)noexcept{return P(*this)>>=x;}
    inline P operator +(const P& x)const noexcept{return P(*this)+=x;}
    inline P operator -(const P& x)const noexcept{return P(*this)-=x;}
    inline P operator -()const noexcept{return P(1,T(0))-=P(*this);}
    inline P operator *(const P& x)const noexcept{return P(*this)*=x;}
    inline P operator /(const P& x)const noexcept{return P(*this)/=x;}
    inline P operator %(const P& x)const noexcept{return P(*this)%=x;}
    bool operator ==(P x){
        for(int i=0;i<(int)max((*this).size(),x.size());++i){
            if(i>=(int)(*this).size()&&x[i]!=T())return 0;
            if(i>=(int)x.size()&&(*this)[i]!=T())return 0;
            if(i<(int)min((*this).size(),x.size()))if((*this)[i]!=x[i])return 0;
        }
        return 1;
    }
    P &operator +=(T x){
        if(this->size()==0)this->resize(1,T(0));
        (*this)[0]+=x;
        return (*this);
    }
    P &operator -=(T x){
        if(this->size()==0)this->resize(1,T(0));
        (*this)[0]-=x;
        return (*this);
    }
    P &operator *=(T x){
        for(int i=0;i<(int)this->size();++i){
            (*this)[i]*=x;
        }
        return (*this);
    }
    P &operator /=(T x){
        if(std::is_same<T,long long>::value){
            for(int i=0;i<(int)this->size();++i){
                (*this)[i]/=x;
            }
            return (*this);
        }
        return (*this)*=(T(1)/x);
    }
    P &operator <<=(int x){
        P ret(x,T(0));
        ret.insert(ret.end(),begin(*this),end(*this));
        return (*this)=ret;
    }
    P &operator >>=(int x){
        if((int)(*this).size()<=x)return (*this)=P();
        P ret;
        ret.insert(ret.end(),begin(*this)+x,end(*this));
        return (*this)=ret;
    }
    P &operator +=(const P& x){
        if(this->size()<x.size())this->resize(x.size(),T(0));
        for(int i=0;i<(int)x.size();++i){
            (*this)[i]+=x[i];
        }
        return (*this);
    }
    P &operator -=(const P& x){
        if(this->size()<x.size())this->resize(x.size(),T(0));
        for(int i=0;i<(int)x.size();++i){
            (*this)[i]-=x[i];
        }
        return (*this);
    }
    P &operator *=(const P& x){
        return (*this)=F()(*this,x);
    }
    P &operator /=(P x){
        if(this->size()<x.size()) {
            this->clear();
            return (*this);
        }
        const int n=this->size()-x.size()+1;
        return (*this) = (rev().pre(n)*x.rev().inv(n)).pre(n).rev(n);
    }
    P &operator %=(const P& x){
        return ((*this)-=(*this)/x*x);
    }
    inline void print(){
        for(int i=0;i<(int)(*this).size();++i)std::cerr<<(*this)[i]<<" \n"[i==(int)(*this).size()-1];
        if((int)(*this).size()==0)std::cerr<<'\n';
    }
    inline P& shrink(){while((*this).back()==0)(*this).pop_back();return (*this);}
    inline P pre(int sz)const{
        return P(begin(*this),begin(*this)+std::min((int)this->size(),sz));
    }
    P rev(int deg=-1){
        P ret(*this);
        if(deg!=-1)ret.resize(deg,T(0));
        reverse(begin(ret),end(ret));
        return ret;
    }
    P inv(int deg=-1){
        assert((*this)[0]!=T(0));
        const int n=deg==-1?this->size():deg;
        P ret({T(1)/(*this)[0]});
        for(int i=1;i<n;i<<=1){
            ret*=(-ret*pre(i<<1)+2).pre(i<<1);
        }
        return ret.pre(n);
    }
    inline P dot(const P& x){
        P ret(*this);
        for(int i=0;i<int(min(this->size(),x.size()));++i){
            ret[i]*=x[i];
        }
        return ret;
    }
    P diff(){
        if((int)(*this).size()<=1)return P();
        P ret(*this);
        for(int i=0;i<(int)ret.size();i++){
            ret[i]*=i;
        }
        return ret>>1;
    }
    P integral(){
        P ret(*this);
        for(int i=0;i<(int)ret.size();i++){
            ret[i]/=i+1;
        }
        return ret<<1;
    }
    P log(int deg=-1){
        assert((*this)[0]==T(1));
        const int n=deg==-1?this->size():deg;
        return (diff()*inv(n)).pre(n-1).integral();
    }
    P exp(int deg=-1){
        assert((*this)[0]==T(0));
        const int n=deg==-1?this->size():deg;
        P ret({T(1)});
        for(int i=1;i<n;i<<=1){
            ret=ret*(pre(i<<1)+1-ret.log(i<<1)).pre(i<<1);
        }
        return ret.pre(n);
    }
    P pow(int c,int deg=-1){
        const int n=deg==-1?this->size():deg;
		long long i=0;
		P ret(*static_cast<P*>(this));
		while(i!=(int)this->size()&&ret[i]==0)i++;
		if(i==(int)this->size())return P(n,0);
		if(i*c>=n)return P(n,0);
		T k=ret[i];
		return ((((ret>>i)/k).log()*c).exp()*(k.pow(c))<<(i*c)).pre(n);
        // const int n=deg==-1?this->size():deg;
        // long long i=0;
        // P ret(*this);
        // while(i!=(int)this->size()&&ret[i]==0)i++;
        // if(i==(int)this->size())return P(n,0);
        // if(i*c>=n)return P(n,0);
        // T k=ret[i];
        // return ((((ret>>i)/k).log()*c).exp()*(k.pow(c))<<(i*c)).pre(n);
        // P x(*this);
        // P ret(1,1);
        // while(c) {
        //     if(c&1){
        //         ret*=x;
        //         if(~deg)ret=ret.pre(deg);
        //     }
        //     x*=x;
        //     if(~deg)x=x.pre(deg);
        //     c>>=1;
        // }
        // return ret;
    }
    P sqrt(int deg=-1){
        const int n=deg==-1?this->size():deg;
        if((*this)[0]==T(0)) {
            for(int i=1;i<(int)this->size();i++) {
                if((*this)[i]!=T(0)) {
                    if(i&1)return{};
                    if(n-i/2<=0)break;
                    auto ret=(*this>>i).sqrt(n-i/2)<<(i/2);
                    if((int)ret.size()<n)ret.resize(n,T(0));
                    return ret;
                }
            }
            return P(n,0);
        }
        P ret({T(1)});
        for(int i=1;i<n;i<<=1){
            ret=(ret+pre(i<<1)*ret.inv(i<<1)).pre(i<<1)/T(2);
        }
        return ret.pre(n);
    }
    P shift(int c){
        const int n=this->size();
        P f(*this),g(n,0);
        for(int i=0;i<n;++i)f[i]*=F().fact(T(i));
        for(int i=0;i<n;++i)g[i]=F().pow(T(c),i)/F().fact(T(i));
        g=g.rev();
        f*=g;
        f>>=n-1;
        for(int i=0;i<n;++i)f[i]/=F().fact(T(i));
        return f;
    }
    T eval(T x){
        T res=0;
        for(int i=(int)this->size()-1;i>=0;--i){
            res*=x;
            res+=(*this)[i];
        }
        return res;
    }
    P mul(const std::vector<std::pair<int,T>>& x){
        int mx=0;
        for(auto [s,t]:x){
            if(mx<s)mx=s;
        }
        P res((int)this->size()+mx);
        for(int i=0;i<(int)this->size();++i){
            for(auto [s,t]:x){
                res[i+s]+=(*this)[i]*t;
            }
        }
        return res;
    }
    P div(const std::vector<std::pair<int,T>>& x){
        P res(*this);
        T cnt=0;
        for(auto [s,t]:x){
            if(s==0)cnt+=t;
        }
        cnt=cnt.inv();
        for(int i=0;i<(int)this->size();++i){
            for(auto [s,t]:x){
                if(s==0)continue;
                if(i>=s)res[i]-=res[i-s]*t*cnt;
            }
        }
        res*=cnt;
        return res;
    }
    static P interpolation(const std::vector<T>&x,const std::vector<T>& y){
        const int n=x.size();
        std::vector<std::pair<P,P>>a(n*2-1);
        std::vector<P> b(n*2-1);
        for(int i=0;i<n;++i)a[i+n-1]=std::make_pair(P{1},P{T()-x[i],1});
        for(int i=n-2;i>=0;--i)a[i]={a[2*i+1].first*a[2*i+2].second+a[2*i+2].first*a[2*i+1].second,a[2*i+1].second*a[2*i+2].second};
        auto d=(a[0].first).multipoint_eval(x);
        for(int i=0;i<n;++i)b[i+n-1]=P{T(y[i]/d[i])};
        for(int i=n-2;i>=0;--i)b[i]=b[2*i+1]*a[2*i+2].second+b[2*i+2]*a[2*i+1].second;
        return b[0];
    }
    static P interpolation(const std::vector<T>& y){
        const int n=y.size();
        std::vector<std::pair<P,P>>a(n*2-1);
        std::vector<P>b(n*2-1);
        for(int i=0;i<n;++i)a[i+n-1]=std::make_pair(P{1},P{T()-i,1});
        for(int i=n-2;i>=0;--i)a[i]={a[2*i+1].first*a[2*i+2].second+a[2*i+2].first*a[2*i+1].second,a[2*i+1].second*a[2*i+2].second};
        for(int i=0;i<n;++i){
            T tmp=F().fact(T(i))*F().pow(T(-1),i)*F().fact(T(n-1-i));
            b[i+n-1]=P{T(y[i]/tmp)};
        }
        for(int i=n-2;i>=0;--i)b[i]=b[2*i+1]*a[2*i+2].second+b[2*i+2]*a[2*i+1].second;
        return b[0];
    }
    std::vector<T> multipoint_eval(const std::vector<T>&x){
        const int n=x.size();
        P* v=new P[2*n-1];
        for(int i=0;i<n;i++)v[i+n-1]={T()-x[i],T(1)};
        for(int i=n-2;i>=0;i--){v[i]=v[i*2+1]*v[i*2+2];}
        v[0]=P(*this)%v[0];v[0].shrink();
        for(int i=1;i<n*2-1;i++){
            v[i]=v[(i-1)/2]%v[i];
            v[i].shrink();
        }
        std::vector<T>res(n);
        for(int i=0;i<n;i++)res[i]=v[i+n-1][0];
        return res;
    }
    P slice(int s,int e,int k){
        P res;
        for(int i=s;i<e;i+=k)res.push_back((*this)[i]);
        return res;
    }
    T nth_term(P q,int64_t x){
        if(x==0)return (*this)[0]/q[0];
        P p(*this);
        P q2=q;
        for(int i=1;i<(int)q2.size();i+=2)q2[i]*=-1;
        q*=q2;
        p*=q2;
        return p.slice(x%2,p.size(),2).nth_term(q.slice(0,q.size(),2),x/2);
    }
    P gcd(P q){
        return *this==P()?q:(q%(*this).shrink()).gcd(*this);
    }
    //(*this)(t(x))
    P manipulate(P t,int deg){
        P s=P(*this);
        if(deg==0)return P();
        if((int)t.size()==1)return P{s.eval(t[0])};
        int k=std::min((int)::sqrt(deg/(::log2(deg)+1))+1,(int)t.size());
        int b=deg/k+1;
        P t2=t.pre(k);
        std::vector<P>table(s.size()/2+1,P{1});
        for(int i=1;i<(int)table.size();i++){
            table[i]=((table[i-1])*t2).pre(deg);
        }
        auto f=[&](auto f,auto l,auto r,int deg)->P{
            if(r-l==1)return P{*l};
            auto m=l+(r-l)/2;
            return f(f,l,m,deg)+(table[m-l]*f(f,m,r,deg)).pre(deg);
        };
        P ans=P();
        P tmp=f(f,s.begin(),s.end(),deg);
        P tmp2=P{1};
        T tmp3=T(1);
        int tmp5=-1;
        P tmp6=t2.diff();
        if(tmp6==P()){
            for(int i=0;i<b;++i){
                if(tmp.size()==0)break;
                ans+=(tmp2*tmp[0]).pre(deg)/tmp3;
                tmp=tmp.diff();
                tmp2=(tmp2*(t-t2)).pre(deg);
                tmp3*=T(i+1);
            }
        }else{
            while(t2[++tmp5]==T());
            P tmp4=(tmp6>>(tmp5-1)).inv(deg);
            for(int i=0;i<b;++i){
                ans+=(tmp*tmp2).pre(deg)/tmp3;
                tmp=((tmp.diff()>>(tmp5-1))*tmp4).pre(deg);
                tmp2=(tmp2*(t-t2)).pre(deg);
                tmp3*=T(i+1);
            }
        }
        return ans;
    }
    //(*this)(t(x))
    P manipulate2(P t,int deg){
        P ans=P();
        P s=(*this).rev();
        for(int i=0;i<(int)s.size();++i){
            ans=(ans*t+s[i]).pre(deg);
        }
        return ans;
    }
    P find_linear_recurrence()const{
        const int n=this->size();
        P b={T(-1)},c={T(-1)};
        T y=T(1);
        for(int i=1;i<=n;++i){
            int l=c.size(),m=b.size();
            T x=0;
            for(int j=0;j<l;++j)x+=c[j]*(*this)[i-l+j];
            b.emplace_back(0);
            m++;
            if(x==T(0))continue;
            T freq=x/y;
            if(l<m){
                auto tmp=c;
                c<<=m-l;
                c-=b*freq;
                b=tmp;
                y=x;
            }else{
                c-=(b*freq)<<(l-m);
            }
        }
        return c;
    }
    static P stirling_second(int n){
        P a(n+1,0),b(n+1,0);
        for(int i=0;i<=n;++i){
            a[i]=F().pow(T(i),n)/F().fact(T(i));
            b[i]=(i%2?T(-1):T(1))/F().fact(T(i));
        }
        return (a*b).pre(n+1);
    }
    void debug(){
        for(int i=0;i<(int)(*this).size();++i)std::cerr<<(*this)[i]<<" \n"[i==(int)(*this).size()-1];
    }
};
#line 3 "math/FPS_mint.hpp"
#include<atcoder/convolution.hpp>
#line 1 "math/ceil_pow2.hpp"
int ceil_pow2(int n) {
    int x = 0;
    while ((1U << x) < (unsigned int)(n)) x++;
    return x;
}
#line 1 "math/mod_pow.hpp"
/**
 * @brief (x^y)%mod
 */

long long mod_pow(long long x,long long y,long long mod){
    long long ret=1;
    while(y>0) {
        if(y&1)(ret*=x)%=mod;
        (x*=x)%=mod;
        y>>=1;
    }
    return ret;
}
#line 4 "math/garner.hpp"

/**
 * 
 * @brief ガーナーのアルゴリズム
 *
 */

long long garner(const std::vector<long long>&a,const std::vector<long long>&mods){
    const int sz=a.size();
    long long coeffs[sz+1]={1,1,1,1};
    long long constants[sz+1]={};
    for(int i=0;i<sz;i++){
        long long v=(mods[i]+a[i]-constants[i])%mods[i]*mod_pow(coeffs[i],mods[i]-2,mods[i])%mods[i];
        for(int j=i+1;j<sz+1;j++) {
            constants[j]=(constants[j]+coeffs[j]*v)%mods[j];
            coeffs[j]=(coeffs[j]*mods[i])%mods[j];
        }
    }
    return constants[sz];
}
#line 6 "math/FPS_mint.hpp"
/**
 * @brief 形式的冪級数(ModInt)
 */

template<typename Mint>
struct _FPS{
    template<typename T>
    T operator()(const T& _s,const T& _t){
        if(_s.size()==0||_t.size()==0)return T();
        const size_t sz=_s.size()+_t.size()-1;
        if((Mint::get_mod()&((1<<ceil_pow2(sz))-1))==1){
            std::vector<atcoder::static_modint<Mint::get_mod()>>s(_s.size()),t(_t.size());
            for(size_t i=0;i<_s.size();++i)s[i]=_s[i].value();
            for(size_t i=0;i<_t.size();++i)t[i]=_t[i].value();
            std::vector<atcoder::static_modint<Mint::get_mod()>> _v=atcoder::convolution(s,t);
            T v(_v.size());
            for (size_t i=0;i<_v.size();++i)v[i]=_v[i].val();
            return v;
        }else{
            std::vector<atcoder::static_modint<1224736769>>s1(_s.size()),t1(_t.size());
            std::vector<atcoder::static_modint<1045430273>>s2(_s.size()),t2(_t.size());
            std::vector<atcoder::static_modint<1007681537>>s3(_s.size()),t3(_t.size());
            for(size_t i=0;i<_s.size();++i){
                s1[i]=_s[i].value();
                s2[i]=_s[i].value();
                s3[i]=_s[i].value();
            }
            for(size_t i=0;i<_t.size();++i){
                t1[i]=_t[i].value();
                t2[i]=_t[i].value();
                t3[i]=_t[i].value();
            }
            auto v1=atcoder::convolution(s1,t1);
            auto v2=atcoder::convolution(s2,t2);
            auto v3=atcoder::convolution(s3,t3);
            T v(sz);
            for(size_t i=0;i<sz;++i){
                v[i]=garner(std::vector<long long>{v1[i].val(),v2[i].val(),v3[i].val()},std::vector<long long>{1224736769,1045430273,1007681537,(long long)Mint::get_mod()});
            }
            return v;
        }
    }
    template<typename T>
    T fact(const T& s){
        return s.fact();
    }
    template<typename T>
    T pow(const T& s,long long i){
        return s.pow(i);
    }
};
template<typename Mint>using fps=FPS_BASE<Mint,_FPS<Mint>>;
#line 2 "math/mod_int_dynamic.hpp"
#include<cstdint>
#line 5 "math/mod_int_dynamic.hpp"

/**
 * @brief ModInt
 */

struct mod_int_dynamic{
    using mint=mod_int_dynamic;
    using u64 = std::uint_fast64_t;
    u64 a;
    mod_int_dynamic(const long long x = 0)noexcept:a(x>=0?x%get_mod():get_mod()-(-x)%get_mod()){}
    u64 &value()noexcept{return a;}
    const u64 &value() const noexcept {return a;}
    mint operator+(const mint rhs)const noexcept{return mint(*this) += rhs;}
    mint operator-(const mint rhs)const noexcept{return mint(*this)-=rhs;}
    mint operator*(const mint rhs) const noexcept {return mint(*this) *= rhs;}
    mint operator/(const mint rhs) const noexcept {return mint(*this) /= rhs;}
    mint &operator+=(const mint rhs) noexcept {
        a += rhs.a;
        if (a >= get_mod())a -= get_mod();
        return *this;
    }
    mint &operator-=(const mint rhs) noexcept {
        if (a<rhs.a)a += get_mod();
        a -= rhs.a;
        return *this;
    }
    mint &operator*=(const mint rhs) noexcept {
        a = a * rhs.a % get_mod();
        return *this;
    }
    mint operator++(int) noexcept {
        a += 1;
        if (a >= get_mod())a -= get_mod();
        return *this;
    }
    mint operator--(int) noexcept {
        if (a<1)a += get_mod();
        a -= 1;
        return *this;
    }
    mint &operator/=(mint rhs) noexcept {
        u64 exp=get_mod()-2;
        while (exp) {
            if (exp % 2) {
                *this *= rhs;
            }
            rhs *= rhs;
            exp /= 2;
        }
        return *this;
    }
    bool operator==(mint x) noexcept {
        return a==x.a;
    }
    bool operator!=(mint x) noexcept {
        return a!=x.a;
    }
    bool operator<(mint x) noexcept {
        return a<x.a;
    }
    bool operator>(mint x) noexcept {
        return a>x.a;
    }
    bool operator<=(mint x) noexcept {
        return a<=x.a;
    }
    bool operator>=(mint x) noexcept {
        return a>=x.a;
    }
    static int root(){
        mint root = 2;
        while(root.pow((get_mod()-1)>>1).a==1)root++;
        return root.a;
    }
    mint pow(long long n)const{
        long long x=a;
        mint ret = 1;
        while(n>0) {
            if(n&1)(ret*=x);
            (x*=x)%=get_mod();
            n>>=1;
        }
        return ret;
    }
    mint inv(){
        return pow(get_mod()-2);
    }
    friend std::ostream& operator<<(std::ostream& lhs, const mint& rhs) noexcept {
        lhs << rhs.a;
        return lhs;
    }
    friend std::istream& operator>>(std::istream& lhs,mint& rhs) noexcept {
        lhs >> rhs.a;
        return lhs;
    }
    constexpr static bool is_static=false;
    static int MOD;
    static u64 get_mod(){
        return MOD;
    }
    static void set_mod(int mod){
        MOD=mod;
    }
};
int mod_int_dynamic::MOD=-1;
#line 4 "math/kth_root.hpp"

template<typename mint>
int kth_root(int n,int k){
    if(k==0){
        if(n==1)return 0;
        else return -1;
    }
    fps<mint>f(k+1,0);
    f[k]=1;
    f[0]=-n;
    random_device rnd;
    for(int times=0;times<10;++times){
        if(f.size()<=2){
            f.resize(k+1);
            f[k]=1;
            f[0]=-n;
        }
        fps<mint>g(k,0),h={1};
        for(int i=0;i<k;++i)g[i]=rnd();
        int t=(mint::get_mod()-1)/2;
        while(t){
            if(t%2)h*=g,h%=f,h.shrink();
            g*=g;
            g%=f;
            g.shrink();
            t/=2;
        }
        f=f.gcd(h+1).shrink();
        if(f.size()==2)return (f[0]/f[1]*(-1)).value();
    }
    return -1;
}