#1016859 (PyPy3) No.2936 Sum of Square of Mex

提出ソース
結果

問題	No.2936 Sum of Square of Mex
ユーザー	hirayuu_yc
提出日時	2024-10-12 13:15:33
言語	PyPy3 (7.3.15)
結果	AC
実行時間	1,150 ms / 2,000 ms
コード長	20,463 bytes
コンパイル時間	196 ms
コンパイル使用メモリ	82,396 KB
実行使用メモリ	247,740 KB
最終ジャッジ日時	2024-10-12 13:16:08
合計ジャッジ時間	31,359 ms
ジャッジサーバーID （参考情報）	judge4 / judge
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	802 ms 224,196 KB
testcase_01	AC	798 ms 224,464 KB
testcase_02	AC	865 ms 226,944 KB
testcase_03	AC	801 ms 224,328 KB
testcase_04	AC	834 ms 224,000 KB
testcase_05	AC	802 ms 224,520 KB
testcase_06	AC	809 ms 223,796 KB
testcase_07	AC	787 ms 224,360 KB
testcase_08	AC	1,104 ms 245,708 KB
testcase_09	AC	1,022 ms 244,688 KB
testcase_10	AC	989 ms 235,276 KB
testcase_11	AC	1,101 ms 246,376 KB
testcase_12	AC	1,000 ms 235,580 KB
testcase_13	AC	1,043 ms 246,032 KB
testcase_14	AC	1,068 ms 245,572 KB
testcase_15	AC	1,089 ms 245,676 KB
testcase_16	AC	1,124 ms 247,488 KB
testcase_17	AC	1,148 ms 247,560 KB
testcase_18	AC	1,107 ms 247,424 KB
testcase_19	AC	1,105 ms 247,204 KB
testcase_20	AC	1,150 ms 247,592 KB
testcase_21	AC	1,111 ms 247,132 KB
testcase_22	AC	1,096 ms 247,740 KB
testcase_23	AC	1,103 ms 247,588 KB
testcase_24	AC	1,085 ms 247,548 KB
testcase_25	AC	1,065 ms 247,052 KB
testcase_26	AC	1,134 ms 246,928 KB
testcase_27	AC	1,143 ms 247,272 KB
testcase_28	AC	1,116 ms 246,940 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
#https://judge.yosupo.jp/submission/82960
class Modulo_Polynominal():
    __slots__=("Poly", "max_degree")

    def __init__(self, Poly=[], max_degree=2*10**5):
        from itertools import zip_longest
        """多項式の定義

        P:係数のリスト
        max_degree

        ※Mod:法はグローバル変数から指定
        """

        if Poly:
            self.Poly=[p%Mod for p in Poly[:max_degree]]
        else:
            self.Poly=[0]
        self.max_degree=max_degree

    def __str__(self):
        return str(self.Poly)

    def __repr__(self):
        return self.__str__()

    def __iter__(self):
        yield from self.Poly

    #=
    def __eq__(self,other):
        from itertools import zip_longest
        return all([a==b for a,b in zip_longest(self.Poly,other.Poly,fillvalue=0)])

    #+,-
    def __pos__(self):
        return self

    def __neg__(self):
        return self.scale(-1)

    #items
    def __getitem__(self, index):
        if isinstance(index, slice):
            return Modulo_Polynominal(self.Poly[index], self.max_degree)
        else:
            if index<0:
                raise IndexError("index is negative (index: {})".format(index))
            elif index>=len(self.Poly):
                return 0
            else:
                return self.Poly[index]

    def __setitem__(self, index, value):
        if index<0:
            raise IndexError("index is negative (index: {})".format(index))
        elif index>=self.max_degree:
            return

        if index>=len(self.Poly):
            self.Poly+=[0]*(index-len(self.Poly)+1)
        self.Poly[index]=value%Mod

    #Boole
    def __bool__(self):
        return any(self.Poly)

    #簡略化
    def reduce(self):
        """ 高次の 0 を切り捨て

        """

        P=self.Poly
        for d in range(len(P)-1,-1,-1):
            if P[d]:
                break
        self.resize(d+1)
        return

    #シフト
    def __lshift__(self,other):
        if other<0:
            return self>>(-other)

        if other>self.max_degree:
            return Modulo_Polynominal([0],self.max_degree)

        G=[0]*other+self.Poly
        return Modulo_Polynominal(G,self.max_degree)

    def __rshift__(self,other):
        if other<0:
            return  self<<(-other)

        return Modulo_Polynominal(self.Poly[other:],self.max_degree)

    #次数
    def degree(self):
        P=self.Poly
        for d in range(len(self.Poly)-1,-1,-1):
            if P[d]:
                return d
        return -float("inf")

    #加法
    def __add__(self,other):
        P=self; Q=other

        if Q.__class__==Modulo_Polynominal:
            N=min(P.max_degree,Q.max_degree)
            A=P.Poly; B=Q.Poly
        else:
            N=P.max_degree
            A=P.Poly; B=Q
        return Modulo_Polynominal(Calc.Add(A,B),N)

    def __radd__(self,other):
        return self+other

    #減法
    def __sub__(self,other):
        P=self; Q=other
        if Q.__class__==Modulo_Polynominal:
            N=min(P.max_degree,Q.max_degree)
            A=P.Poly; B=Q.Poly
        else:
            N=P.max_degree
            A=P.Poly; B=Q
        return Modulo_Polynominal(Calc.Sub(A,B),N)

    def __rsub__(self,other):
        return (-self)+other

    #乗法
    def __mul__(self,other):
        P=self
        Q=other
        if Q.__class__==Modulo_Polynominal:
            a=b=0
            for x in P.Poly:
                if x:
                    a+=1
            for y in Q.Poly:
                if y:
                    b+=1

            if a>b:
                P,Q=Q,P

            P.reduce();Q.reduce()
            U,V=P.Poly,Q.Poly
            M=min(P.max_degree,Q.max_degree)
            if a<2*P.max_degree.bit_length():
                B=[0]*(len(U)+len(V)-1)
                for i in range(len(U)):
                    if U[i]:
                        for j in range(len(V)):
                            B[i+j]+=U[i]*V[j]
                            if B[i+j]>Mod:
                                B[i+j]-=Mod
            else:
                B=Calc.Convolution(U,V)[:M]
            B=Modulo_Polynominal(B,M)
            B.reduce()
            return B
        else:
            return self.scale(other)

    def __rmul__(self,other):
        return self.scale(other)

    #除法
    def __floordiv__(self,other):
        if not other:
            raise ZeroDivisionError
        if isinstance(other,int):
            return self/other

        self.reduce()
        other.reduce()

        return Modulo_Polynominal(Calc.Floor_Div(self.Poly, other.Poly),
                                max(self.max_degree, other.max_degree))

    def __rfloordiv__(self,other):
        if not self:
            raise ZeroDivisionError

        if isinstance(other,int):
            return Modulo_Polynominal([],self.max_degree)

    #剰余
    def __mod__(self,other):
        if not other:
            return ZeroDivisionError
        self.reduce(); other.reduce()
        r=Modulo_Polynominal(Calc.Mod(self.Poly, other.Poly),
                            max(self.max_degree, other.max_degree))
        r.reduce()
        return r

    def __rmod__(self,other):
        if not self:
            raise ZeroDivisionError
        r=other-(other//self)*self
        r.reduce()
        return r

    def __divmod__(self,other):
        q=self//other
        r=self-q*other; r.reduce()
        return (q,r)

    #累乗
    def __pow__(self,other):
        if other.__class__==int:
            n=other
            m=abs(n)

            Q=self
            A=Modulo_Polynominal([1],self.max_degree)
            while m>0:
                if m&1:
                    A*=Q
                m>>=1
                Q*=Q

            if n>=0:
                return A
            else:
                return A.inverse()
        else:
            P=Log(self)
            return Exp(P*other)

    #逆元
    def inverse(self, deg=None):
        assert self.Poly[0],"定数項が0"

        if deg==None:
            deg=self.max_degree

        return Modulo_Polynominal(Calc.Inverse(self.Poly, deg), self.max_degree)

    #除法
    def __truediv__(self,other):
        if isinstance(other, Modulo_Polynominal):
            return self*other.inverse()
        else:
            return pow(other,Mod-2,Mod)*self

    def __rtruediv__(self,other):
        return other*self.inverse()

    #スカラー倍
    def scale(self, s):
        return Modulo_Polynominal(Calc.Times(self.Poly,s),self.max_degree)

    #最高次の係数
    def leading_coefficient(self):
        for x in self.Poly[::-1]:
            if x:
                return x
        return 0

    def censor(self, N=-1, Return=False):
        """ N 次以上の係数をカット
        """

        if N==-1:
            N=self.max_degree

        N=min(N, self.max_degree)

        if Return:
            return Modulo_Polynominal(self.Poly[:N],self.max_degree)
        else:
            self.Poly=self.Poly[:N]

    def resize(self, N, Return=False):
        """ 強制的に Poly の配列の長さを N にする.

        """

        N=min(N, self.max_degree)
        P=self
        if Return:
            if len(P.Poly)>N:
                E=P.Poly[:N]
            else:
                E=P.Poly+[0]*(N-len(P.Poly))
            return Modulo_Polynominal(E,P.max_degree)
        else:
            if len(P.Poly)>N:
                del P.Poly[N:]
            else:
                P.Poly+=[0]*(N-len(P.Poly))

    #代入
    def substitution(self, a):
        """ a を (形式的に) 代入した値を求める.

        a: int
        """

        y=0
        t=1
        for p in self.Poly:
            y=(y+p*t)%Mod
            t=(t*a)%Mod
        return y

class Calculator:
    def __init__(self):
        self.primitive=self.__primitive_root()
        self.__build_up()

    def __primitive_root(self):
        p=Mod
        if p==2:
            return 1
        if p==998244353:
            return 3
        if p==10**9+7:
            return 5
        if p==163577857:
            return 23
        if p==167772161:
            return 3
        if  p==469762049:
            return 3

        fac=[]
        q=2
        v=p-1

        while v>=q*q:
            e=0
            while v%q==0:
                e+=1
                v//=q

            if e>0:
                fac.append(q)
            q+=1

        if v>1:
            fac.append(v)

        g=2
        while g<p:
            if pow(g,p-1,p)!=1:
                return None

            flag=True
            for q in fac:
                if pow(g,(p-1)//q,p)==1:
                    flag=False
                    break

            if flag:
                return g

            g+=1

    #参考元: https://judge.yosupo.jp/submission/72676
    def __build_up(self):
        rank2=(~(Mod-1)&(Mod-2)).bit_length()
        root=[0]*(rank2+1); iroot=[0]*(rank2+1)
        rate2=[0]*max(0, rank2-1); irate2=[0]*max(0, rank2-1)
        rate3=[0]*max(0, rank2-2); irate3=[0]*max(0, rank2-2)

        root[-1]=pow(self.primitive, (Mod-1)>>rank2, Mod)
        iroot[-1]=pow(root[-1], Mod-2, Mod)

        for i in range(rank2)[::-1]:
            root[i]=root[i+1]*root[i+1]%Mod
            iroot[i]=iroot[i+1]*iroot[i+1]%Mod

        prod=iprod=1
        for i in range(rank2-1):
            rate2[i]=root[i+2]*prod%Mod
            irate2[i]=iroot[i+2]*prod%Mod
            prod*=iroot[i+2]; prod%=Mod
            iprod*=root[i+2]; iprod%=Mod

        prod=iprod = 1
        for i in range(rank2-2):
            rate3[i]=root[i + 3]*prod%Mod
            irate3[i]=iroot[i + 3]*iprod%Mod
            prod*=iroot[i + 3]; prod%=Mod
            iprod*=root[i + 3]; iprod%=Mod

        self.root=root; self.iroot=iroot
        self.rate2=rate2; self.irate2=irate2
        self.rate3=rate3; self.irate3=irate3

    def Add(self, A,B):
        """ 必要ならば末尾に元を追加して, [A[i]+B[i]] を求める.

        """
        if type(A)==int:
            A=[A]

        if type(B)==int:
            B=[B]

        m=min(len(A), len(B))
        C=[0]*m
        for i in range(m):
            C[i]=(A[i]+B[i])%Mod
        C.extend(A[m:])
        C.extend(B[m:])
        return C

    def Sub(self, A,B):
        """ 必要ならば末尾に元を追加して, [A[i]-B[i]] を求める.

        """
        if type(A)==int:
            A=[A]

        if type(B)==int:
            B=[B]

        m=min(len(A), len(B))
        C=[0]*m
        for i in range(m):
            C[i]=(A[i]-B[i])%Mod
        C.extend(A[m:])
        C.extend([-b%Mod for b in B[m:]])
        return C

    def Times(self,A,k):
        """ [k*A[i]] を求める.

        """
        return [k*a%Mod for a in A]

    #参考元 https://judge.yosupo.jp/submission/72676
    def NTT(self,A):
        """A に Mod を法とする数論変換を施す

        ※ Mod はグローバル変数から指定

        References:
        https://github.com/atcoder/ac-library/blob/master/atcoder/convolution.hpp
        https://judge.yosupo.jp/submission/72676
        """

        N=len(A)
        H=(N-1).bit_length()
        l=0

        I=self.root[2]
        rate2=self.rate2; rate3=self.rate3

        while l<H:
            if H-l==1:
                p=1<<(H-l-1)
                rot=1
                for s in range(1<<l):
                    offset=s<<(H-l)
                    for i in range(p):
                        x=A[i+offset]; y=A[i+offset+p]*rot%Mod
                        A[i+offset]=(x+y)%Mod
                        A[i+offset+p]=(x-y)%Mod

                    if s+1!=1<<l:
                        rot*=rate2[(~s&-~s).bit_length()-1]
                        rot%=Mod
                l+=1
            else:
                p=1<<(H-l-2)
                rot=1
                for s in range(1<<l):
                    rot2=rot*rot%Mod
                    rot3=rot2*rot%Mod
                    offset=s<<(H-l)
                    for i in range(p):
                        a0=A[i+offset]
                        a1=A[i+offset+p]*rot
                        a2=A[i+offset+2*p]*rot2
                        a3=A[i+offset+3*p]*rot3

                        alpha=(a1-a3)%Mod*I

                        A[i+offset]=(a0+a2+a1+a3)%Mod
                        A[i+offset+p]=(a0+a2-a1-a3)%Mod
                        A[i+offset+2*p]=(a0-a2+alpha)%Mod
                        A[i+offset+3*p]=(a0-a2-alpha)%Mod

                    if s+1!=1<<l:
                        rot*=rate3[(~s&-~s).bit_length()-1]
                        rot%=Mod
                l+=2

    #参考元 https://judge.yosupo.jp/submission/72676
    def Inverse_NTT(self, A):
        """A を Mod を法とする逆数論変換を施す

        ※ Mod はグローバル変数から指定

        References:
        https://github.com/atcoder/ac-library/blob/master/atcoder/convolution.hpp
        https://judge.yosupo.jp/submission/72676
        """
        N=len(A)
        H=(N-1).bit_length()
        l=H

        J=self.iroot[2]
        irate2=self.rate2; irate3=self.irate3

        while l:
            if l==1:
                p=1<<(H-l)
                irot=1
                for s in range(1<<(l-1)):
                    offset=s<<(H-l+1)
                    for i in range(p):
                        x=A[i+offset]; y=A[i+offset+p]
                        A[i+offset]=(x+y)%Mod
                        A[i+offset+p]=(x-y)*irot%Mod

                    if s+1!=1<<(l-1):
                        irot*=irate2[(~s&-~s).bit_length()-1]
                        irot%=Mod
                l-=1
            else:
                p=1<<(H-l)
                irot=1
                for s in range(1<<(l-2)):
                    irot2=irot*irot%Mod
                    irot3=irot2*irot%Mod
                    offset=s<<(H-l+2)
                    for i in range(p):
                        a0=A[i+offset]
                        a1=A[i+offset+p]
                        a2=A[i+offset+2*p]
                        a3=A[i+offset+3*p]

                        beta=(a2-a3)*J%Mod

                        A[i+offset]=(a0+a1+a2+a3)%Mod
                        A[i+offset+p]=(a0-a1+beta)*irot%Mod
                        A[i+offset+2*p]=(a0+a1-a2-a3)*irot2%Mod
                        A[i+offset+3*p]=(a0-a1-beta)*irot3%Mod

                    if s+1!=1<<(l-2):
                        irot*=irate3[(~s&-~s).bit_length()-1]
                        irot%=Mod
                l-=2
        N_inv=pow(N,Mod-2,Mod)
        for i in range(N):
            A[i]=N_inv*A[i]%Mod

    def Convolution(self, A, B):
        """ A, B で Mod を法とする畳み込みを求める.

        ※ Mod はグローバル変数から指定
        """
        if not A or not B:
            return []

        N=len(A)
        M=len(B)
        L=M+N-1

        if min(N,M)<=50:
            if N<M:
                N,M=M,N
                A,B=B,A
            C=[0]*L
            for i in range(N):
                for j in range(M):
                    C[i+j]+=A[i]*B[j]
                    C[i+j]%=Mod

            return C

        H=L.bit_length()
        K=1<<H

        A=A+[0]*(K-N)
        B=B+[0]*(K-M)

        self.NTT(A)
        self.NTT(B)

        for i in range(K):
            A[i]=A[i]*B[i]%Mod

        self.Inverse_NTT(A)

        return A[:L]

    def Autocorrelation(self, A):
        """ A 自身に対して,Mod を法とする畳み込みを求める.

        ※ Mod はグローバル変数から指定
        """
        N=len(A)
        L=2*N-1

        if N<=50:
            C=[0]*L
            for i in range(N):
                for j in range(N):
                    C[i+j]+=A[i]*A[j]
                    C[i+j]%=Mod
            return C

        H=L.bit_length()
        K=1<<H

        A=A+[0]*(K-N)

        self.NTT(A)

        for i in range(K):
            A[i]=A[i]*A[i]%Mod
        self.Inverse_NTT(A)

        return A[:L]

    def Multiple_Convolution(self, *A):
        """ A=(A[0], A[1], ..., A[d-1]) で Mod を法とする畳み込みを行う.

        """
        from heapq import heapify,heappush,heappop
        Q=[(len(a), a) for a in A]
        heapify(Q)
        while len(Q)>1:
            m,a=heappop(Q)
            n,b=heappop(Q)
            heappush(Q, (m+n, self.Convolution(a,b)))
        return Q[0][1]

    def Inverse(self, F, length=None):
        if length==None:
            M=len(F)
        else:
            M=length

        if len(F)<=M.bit_length():
            """
            愚直に漸化式を用いて求める.
            計算量:Pの次数をK, 求めたい項の個数をNとして, O(NK)
            """
            c=F[0]
            c_inv=pow(c,Mod-2,Mod)

            N=len(F)
            R=[-c_inv*a%Mod for a in F[1:]][::-1]
            G=[0]*M
            G[0]=1
            Q=[0]*(N-2)+[1]

            for k in range(1,M):
                a=0
                for x,y in zip(Q,R):
                    a+=x*y
                a%=Mod
                G[k]=a
                Q.append(a)
                Q=Q[1:]

            G=[c_inv*g%Mod for g in G]
        else:
            """
            FFTの理論を応用して求める.
            計算量: 求めたい項の個数をNとして, O(N log N)

            Reference: https://judge.yosupo.jp/submission/42413
            """

            N=len(F)
            r=pow(F[0],Mod-2,Mod)

            m=1
            G=[r]
            while m<M:
                A=F[:min(N, 2*m)]; A+=[0]*(2*m-len(A))
                B=G.copy(); B+=[0]*(2*m-len(B))

                Calc.NTT(A); Calc.NTT(B)
                for i in range(2*m):
                    A[i]=A[i]*B[i]%Mod

                Calc.Inverse_NTT(A)
                A=A[m:]+[0]*m
                Calc.NTT(A)
                for i in range(2*m):
                    A[i]=-A[i]*B[i]%Mod
                Calc.Inverse_NTT(A)

                G.extend(A[:m])
                m<<=1
            G=G[:M]
        return G

    def Floor_Div(self, F,G):
        assert F[-1]
        assert G[-1]

        F_deg=len(F)-1
        G_deg=len(G)-1

        if F_deg<G_deg:
            return []

        m=F_deg-G_deg+1
        return self.Convolution(F[::-1], Calc.Inverse(G[::-1],m))[m-1::-1]

    def Mod(self, F,G):
        while F and F[-1]==0: F.pop()
        while G and G[-1]==0: G.pop()

        if not F:
            return []

        return Calc.Sub(F, Calc.Convolution(Calc.Floor_Div(F,G),G))

def Taylor_Shift(P,a):
    """与えられた多項式 P に対して, P(X+a) を求める.

    P: Polynominal
    a: int
    """

    N=len(P.Poly)-1

    fact=[0]*(N+1)
    fact[0]=1
    for i in range(1,N+1):
        fact[i]=(fact[i-1]*i)%Mod

    fact_inv=[0]*(N+1)
    fact_inv[-1]=pow(fact[-1],Mod-2,Mod)

    for i in range(N-1,-1,-1):
        fact_inv[i]=(fact_inv[i+1]*(i+1))%Mod

    F=P.Poly.copy()
    for i in range(N+1):
        F[i]=(F[i]*fact[i])%Mod

    G=[0]*(N+1)
    c=1
    for i in range(N+1):
        G[i]=(c*fact_inv[i])%Mod
        c=(c*a)%Mod
    G.reverse()

    H=Calc.Convolution(F,G)[N:]
    for i in range(len(H)):
        H[i]=(H[i]*fact_inv[i])%Mod

    return Modulo_Polynominal(H,P.max_degree)

def Stirling_2nd(N):
    """ k=0,1, ..., N に対する第 II 種 Stirling 数を求める.

    """

    fact=[0]*(N+1); fact[0]=1
    for i in range(1,N+1):
        fact[i]=i*fact[i-1]%Mod

    fact_inv=[0]*(N+1); fact_inv[-1]=pow(fact[-1],Mod-2,Mod)
    for i in range(N-1,-1,-1):
        fact_inv[i]=(i+1)*fact_inv[i+1]%Mod

    A=[pow(i,N,Mod)*fact_inv[i]%Mod for i in range(N+1)]
    B=[fact_inv[i] if i&1==0 else -fact_inv[i] for i in range(N+1)]
    return Calc.Convolution(A,B)[:N+1]

#==================================================
fact=[]
rev=[]
def comb_init(N=1000000,MOD=998244353):
    global fact,rev
    fact=[1]
    rev=[1]
    for i in range(N):
        fact.append((fact[-1]*(i+1))%MOD)
        rev.append((rev[-1]*pow(i+1,-1,MOD))%MOD)


Mod=998244353; Calc=Calculator(); comb_init()
N,M=map(int,input().split())
cf=[1]*1000001
for i in range(1000000):
    cf[i+1]=cf[i]*(M+1-i)
    cf[i+1]%=998244353
def comb(n,r):
    r=n-r
    if r<0:
        return 0
    return (cf[r]*rev[r])%998244353
stir=Stirling_2nd(N)
ans=0
for i in range(N+1):
    if i>M+1:
        break
    stir[i]*=fact[i]
    ans+=stir[i]*(i*i*comb(M+1,M-i+1)-(i*2-1)*(M-i+1)*comb(M+1,M-i+2)+(M-i+1)*(M-i+2)*comb(M+1,M-i+3))
    ans%=Mod
print(ans)
yukicoder

結果

テストケース

ソースコード
