#747044 (PyPy3) No.1195 数え上げを愛したい（文字列編）

提出ソース

結果

問題	No.1195 数え上げを愛したい（文字列編）
ユーザー	👑 Kazun
提出日時	2022-03-19 04:13:23
言語	PyPy3 (7.3.15)
結果	AC
実行時間	833 ms / 3,000 ms
コード長	7,907 bytes
コンパイル時間	234 ms
コンパイル使用メモリ	82,576 KB
実行使用メモリ	127,840 KB
最終ジャッジ日時	2024-04-14 17:15:54
合計ジャッジ時間	15,230 ms
ジャッジサーバーID （参考情報）	judge2 / judge5

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

テストケース

テストケース表示

入力	結果	実行時間実行使用メモリ
testcase_00	AC	831 ms 127,752 KB
testcase_01	AC	831 ms 127,596 KB
testcase_02	AC	833 ms 127,628 KB
testcase_03	AC	303 ms 91,056 KB
testcase_04	AC	338 ms 91,488 KB
testcase_05	AC	323 ms 76,100 KB
testcase_06	AC	41 ms 54,424 KB
testcase_07	AC	42 ms 54,156 KB
testcase_08	AC	258 ms 85,584 KB
testcase_09	AC	831 ms 127,840 KB
testcase_10	AC	480 ms 101,456 KB
testcase_11	AC	799 ms 125,680 KB
testcase_12	AC	777 ms 124,396 KB
testcase_13	AC	658 ms 110,184 KB
testcase_14	AC	451 ms 101,628 KB
testcase_15	AC	477 ms 101,184 KB
testcase_16	AC	462 ms 100,460 KB
testcase_17	AC	263 ms 85,680 KB
testcase_18	AC	777 ms 124,200 KB
testcase_19	AC	780 ms 124,080 KB
testcase_20	AC	663 ms 110,088 KB
testcase_21	AC	794 ms 124,996 KB
testcase_22	AC	532 ms 104,348 KB
testcase_23	AC	41 ms 54,312 KB
testcase_24	AC	42 ms 54,888 KB
testcase_25	AC	41 ms 54,956 KB

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

ソースコード

raw source code

class Convoluter:
    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

    #参考元 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=N+M-1
        if min(N,M)<=20:
            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<=20:
            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]
#==================================================
from heapq import heapify,heappop,heappush

S=input()
N=len(S)

chi=[0]*26
for a in S:
    chi[ord(a)-ord("a")]+=1

Mod=998244353; Conv=Convoluter()
Fact=[1]*(N+1)
Fact_inv=[1]*(N+1)
for i in range(1,N+1):
    Fact[i]=i*Fact[i-1]%Mod
    Fact_inv[i]=pow(Fact[i],Mod-2,Mod)

Q=[(chi[c],Fact_inv[:chi[c]+1]) for c in range(26)]
heapify(Q)

while len(Q)>=2:
    a,A=heappop(Q)
    b,B=heappop(Q)
    heappush(Q,(a+b,Conv.Convolution(A,B)))

_,A=heappop(Q)
X=0
for d,a in enumerate(A):
    if d>=1:
        X+=Fact[d]*a
        X%=Mod

print(X)

yukicoder

結果

テストケース

ソースコード