#840335 (PyPy3) No.1516 simple 門松列 problem Re:MASTER

提出ソース
結果

問題	No.1516 simple 門松列 problem Re:MASTER
ユーザー	vwxyz
提出日時	2023-02-16 07:15:47
言語	PyPy3 (7.3.15)
結果	AC
実行時間	2,674 ms / 6,000 ms
コード長	23,099 bytes
コンパイル時間	180 ms
コンパイル使用メモリ	82,304 KB
実行使用メモリ	80,548 KB
最終ジャッジ日時	2024-07-18 10:22:46
合計ジャッジ時間	12,571 ms
ジャッジサーバーID （参考情報）	judge1 / judge3
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	38 ms 58,624 KB
testcase_01	AC	209 ms 77,056 KB
testcase_02	AC	1,118 ms 78,592 KB
testcase_03	AC	47 ms 66,432 KB
testcase_04	AC	46 ms 66,688 KB
testcase_05	AC	51 ms 67,840 KB
testcase_06	AC	88 ms 72,832 KB
testcase_07	AC	215 ms 76,416 KB
testcase_08	AC	293 ms 76,544 KB
testcase_09	AC	52 ms 65,920 KB
testcase_10	AC	89 ms 71,528 KB
testcase_11	AC	52 ms 67,200 KB
testcase_12	AC	47 ms 66,560 KB
testcase_13	AC	45 ms 66,304 KB
testcase_14	AC	2,238 ms 79,872 KB
testcase_15	AC	1,035 ms 77,568 KB
testcase_16	AC	510 ms 76,480 KB
testcase_17	AC	211 ms 76,416 KB
testcase_18	AC	94 ms 76,232 KB
testcase_19	AC	63 ms 71,296 KB
testcase_20	AC	2,674 ms 80,548 KB
testcase_21	AC	2,549 ms 80,128 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
import sys
readline=sys.stdin.readline

class Matrix:
    def __init__(self,H=0,W=0,matrix=False,eps=0,mod=0,identity=0):
        if identity:
            if H:
                self.H=H
                self.W=H
            else:
                self.H=W
                self.W=W
            self.matrix=[[0]*self.W for i in range(self.H)]
            for i in range(self.H):
                self.matrix[i][i]=identity
        elif matrix:
            self.matrix=matrix
            self.H=len(self.matrix)
            self.W=len(self.matrix[0]) if self.matrix else 0
        else:
            self.H=H
            self.W=W
            self.matrix=[[0]*self.W for i in range(self.H)]
        mod=mod
        self.eps=eps

    def __eq__(self,other):
        if type(other)!=Matrix:
            return False
        if self.H!=other.H:
            return False
        if mod:
            for i in range(self.H):
                for j in range(self.W):
                    if self.matrix[i][j]%mod!=other.matrix[i][j]%mod:
                        return False
        else:
            for i in range(self.H):
                for j in range(self.W):
                    if self.eps<abs(self.matrix[i][j]-other.matrix[i][j]):
                        return False
        return True

    def __ne__(self,other):
        if type(other)!=Matrix:
            return True
        if self.H!=other.H:
            return True
        if mod:
            for i in range(self.H):
                for j in range(self.W):
                    if self.matrix[i][j]%mod!=other.matrix[i][j]%mod:
                        return True
        else:
            for i in range(self.H):
                for j in range(self.W):
                    if self.eps<abs(self.matrix[i][j]-other.matrix[i][j]):
                        return True
        return False

    def __add__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            if mod:
                summ=Matrix(matrix=[[(self.matrix[i][j]+other.matrix[i][j])%mod for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
            else:
                summ=Matrix(matrix=[[self.matrix[i][j]+other.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            if mod:
                summ=Matrix(matrix=[[(self.matrix[i][j]+other)%mod for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
            else:
                summ=Matrix(matrix=[[self.matrix[i][j]+other for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return summ

    def __sub__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            if mod:
                diff=Matrix(matrix=[[(self.matrix[i][j]-other.matrix[i][j])%mod for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
            else:
                diff=Matrix(matrix=[[self.matrix[i][j]-other.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            if mod:
                diff=Matrix(matrix=[[(self.matrix[i][j]-other)%mod for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
            else:
                diff=Matrix(matrix=[[self.matrix[i][j]-other for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return diff

    def __mul__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            if mod:
                prod=Matrix(matrix=[[(self.matrix[i][j]*other.matrix[i][j])%mod for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
            else:
                prod=Matrix(matrix=[[self.matrix[i][j]*other.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            if mod:
                prod=Matrix(matrix=[[(self.matrix[i][j]*other)%mod for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
            else:
                prod=Matrix(matrix=[[self.matrix[i][j]*other for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return prod

    def __matmul__(self,other):
        if type(other)==Matrix:
            assert self.W==other.H
            prod=Matrix(H=self.H,W=other.W,eps=self.eps,mod=mod)
            for i in range(self.H):
                for j in range(other.W):
                    for k in range(self.W):
                        prod.matrix[i][j]+=self.matrix[i][k]*other.matrix[k][j]
                        if mod:
                            prod.matrix[i][j]%=mod
        elif type(other)==int:
            assert self.H==self.W
            if other==0:
                prod=Matrix(H=self.H,eps=self.eps,mod=mod,identity=1)
            elif other==1:
                prod=Matrix(matrix=[[self.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
            else:
                prod=Matrix(H=self.H,eps=self.eps,mod=mod,identity=1)
                doub=Matrix(matrix=[[self.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
                while other>=2:
                    if other&1:
                        prod@=doub
                    doub@=doub
                    other>>=1
                prod@=doub
        return prod

    def __truediv__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            if mod:
                quot=Matrix(matrix=[[(self.matrix[i][j]*MOD(mod).Pow(other.matrix[i][j],-1))%mod for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
            else:
                quot=Matrix(matrix=[[self.matrix[i][j]/other.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            if mod:
                inve=MOD(mod).Pow(other,-1)
                quot=Matrix(matrix=[[(self.matrix[i][j]*inve)%mod for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
            else:
                quot=Matrix(matrix=[[self.matrix[i][j]/other for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return quot

    def __floordiv__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            quot=Matrix(matrix=[[self.matrix[i][j]//other.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            quot=Matrix(matrix=[[self.matrix[i][j]//other for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return quot

    def __mod__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            rema=Matrix(matrix=[[self.matrix[i][j]%other.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            rema=Matrix(matrix=[[self.matrix[i][j]%other for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return rema

    def __pow__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            if mod:
                powe=Matrix(matrix=[[pow(self.matrix[i][j],other.matrix[i][j],mod) for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
            else:
                powe=Matrix(matrix=[[pow(self.matrix[i][j],other.matrix[i][j]) for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            if mod:
                powe=Matrix(matrix=[[pow(self.matrix[i][j],other,mod) for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
            else:
                powe=Matrix(matrix=[[pow(self.matrix[i][j],other) for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return powe

    def __lshift__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            lshi=Matrix(matrix=[[self.matrix[i][j]<<other.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            lshi=Matrix(matrix=[[self.matrix[i][j]<<other for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return lshi

    def __rshift__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            rshi=Matrix(matrix=[[self.matrix[i][j]>>other.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            rshi=Matrix(matrix=[[self.matrix[i][j]>>other for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return rshi

    def __and__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            conj=Matrix(matrix=[[self.matrix[i][j]&other.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            conj=Matrix(matrix=[[self.matrix[i][j]&other for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return conj

    def __or__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            disj=Matrix(matrix=[[self.matrix[i][j]|other.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            disj=Matrix(matrix=[[self.matrix[i][j]|other for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return disj

    def __xor__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            excl=Matrix(matrix=[[self.matrix[i][j]^other.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            excl=Matrix(matrix=[[self.matrix[i][j]^other for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return excl

    def __iadd__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]+=other.matrix[i][j]
                    if mod:
                        self.matrix[i][j]%=mod
        else:
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]+=other
                    if mod:
                        self.matrix[i][j]%=mod
        return self

    def __isub__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]-=other.matrix[i][j]
                    if mod:
                        self.matrix[i][j]%=mod
        else:
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]-=other
                    if mod:
                        self.matrix[i][j]%=mod
        return self

    def __imul__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]*=other.matrix[i][j]
                    if mod:
                        self.matrix[i][j]%=mod
        else:
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]*=other
                    if mod:
                        self.matrix[i][j]%=mod
        return self

    def __imatmul__(self,other):
        if type(other)==Matrix:
            assert self.W==other.H
            prod=Matrix(H=self.H,W=other.W,eps=self.eps,mod=mod)
            for i in range(self.H):
                for j in range(other.W):
                    for k in range(self.W):
                        prod.matrix[i][j]+=self.matrix[i][k]*other.matrix[k][j]
                        if mod:
                            prod.matrix[i][j]%=mod
        elif type(other)==int:
            assert self.H==self.W
            if other==0:
                return Matrix(H=self.H,eps=self.eps,mod=mod,identity=1)
            elif other==1:
                prod=Matrix(matrix=[[self.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
            else:
                prod=Matrix(H=self.H,eps=self.eps,mod=mod,identity=1)
                doub=self
                while other>=2:
                    if other&1:
                        prod@=doub
                    doub@=doub
                    other>>=1
                prod@=doub
        return prod

    def __itruediv__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            for i in range(self.H):
                for j in range(self.W):
                    if mod:
                        self.matrix[i][j]=self.matrix[i][j]*MOD(mod).Pow(other.matrix[i][j],-1)%mod
                    else:
                        self.matrix[i][j]/=other.matrix[i][j]
        else:
            if mod:
                inve=MOD(mod).Pow(other,-1)
            for i in range(self.H):
                for j in range(self.W):
                    if mod:
                        self.matrix[i][j]=self.matrix[i][j]*inve%mod
                    else:
                        self.matrix[i][j]/=other
        return self

    def __ifloordiv__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]//=other.matrix[i][j]
        else:
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]//=other
        return self

    def __imod__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]%=other.matrix[i][j]
        else:
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]%=other
        return self

    def __ipow__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            for i in range(self.H):
                for j in range(self.W):
                    if mod:
                        self.matrix[i][j]=pow(self.matrix[i][j],other.matrix[i][j],mod)
                    else:
                        self.matrix[i][j]=pow(self.matrix[i][j],other.matrix[i][j])
        else:
            for i in range(self.H):
                for j in range(self.W):
                    if mod:
                        self.matrix[i][j]=pow(self.matrix[i][j],other,mod)
                    else:
                        self.matrix[i][j]=pow(self.matrix[i][j],other)
        return self

    def __ilshift__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]<<=other.matrix[i][j]
        else:
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]<<=other
        return self

    def __irshift__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]>>=other.matrix[i][j]
        else:
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]>>=other
        return self

    def __iand__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]&=other.matrix[i][j]
        else:
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]&=other
        return self

    def __ior__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]|=other.matrix[i][j]
        else:
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]|=other
        return self

    def __ixor__(self,other):
        if type(other)==Matrix:
            assert self.H==other.H
            assert self.W==other.W
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]^=other.matrix[i][j]
        else:
            for i in range(self.H):
                for j in range(self.W):
                    self.matrix[i][j]^=other
        return self

    def __neg__(self):
        if mod:
            nega=Matrix(matrix=[[(-self.matrix[i][j])%mod for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        else:
            nega=Matrix(matrix=[[-self.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return nega

    def __pos__(self):
        posi=Matrix(matrix=[[self.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return posi

    def __invert__(self):
        inve=Matrix(matrix=[[~self.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return inve

    def __abs__(self):
        abso=Matrix(matrix=[[abs(self.matrix[i][j]) for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        return abso

    def __getitem__(self,i):
        if type(i)==int:
            return self.matrix[i]
        elif type(i)==tuple:
            i,j=i
            if type(i)==int:
                i=slice(i,i+1)
            if type(j)==int:
                j=slice(j,j+1)
            return Matrix(matrix=[lst[j] for lst in self.matrix[i]],eps=self.eps,mod=mod)

    def __contains__(self,x):
        for i in range(self.H):
            if x in self.matrix[i]:
                return True
        return False

    def __str__(self):
        digit=[max(len(str(self.matrix[i][j])) for i in range(self.H)) for j in range(self.W)]
        return "\n".join([(" [" if i else "[[")+", ".join([str(self.matrix[i][j]).rjust(digit[j]," ") for j in range(self.W)])+"]" for i in range(self.H)])+"]"

    def __bool__(self):
        return True

    def Transpose(self):
        return Matrix(matrix=[[self.matrix[i][j] for i in range(self.H)] for j in range(self.W)])

    def Trace(self):
        assert self.H==self.W
        trace=sum(self.matrix[i][i] for i in range(self.H))
        if mod:
            trace%=mod
        return trace

    def Elem_Raw_Operate_1(self,i0,i1):
        self.matrix[i0],self.matrix[i1]=self.matrix[i1],self.matrix[i0]

    def Elem_Raw_Operate_2(self,i,c):
        if mod:
            self.matrix[i]=[self.matrix[i][j]*c%mod for j in range(self.W)]
        else:
            self.matrix[i]=[self.matrix[i][j]*c for j in range(self.W)]

    def Elem_Raw_Operate_3(self,i0,i1,c):
        if mod:
            self.matrix[i0]=[(self.matrix[i0][j]+c*self.matrix[i1][j])%mod for j in range(self.W)]
        else:
            self.matrix[i0]=[self.matrix[i0][j]+c*self.matrix[i1][j] for j in range(self.W)]

    def Elimination(self,determinant=False,inverse_matrix=False,linear_equation=False,rank=False,upper_triangular=False):
        h=0
        ut=Matrix(matrix=[[self.matrix[i][j] for j in range(self.W)] for i in range(self.H)],eps=self.eps,mod=mod)
        if determinant or inverse_matrix:
            assert self.H==self.W
            det=1
        if inverse_matrix:
            assert self.H==self.W
            im=Matrix(H=self.H,eps=self.eps,mod=mod,identity=1)
        if linear_equation:
            assert self.H==linear_equation.H
            le=Matrix(matrix=[[linear_equation.matrix[i][j] for j in range(linear_equation.W)] for i in range(linear_equation.H)],eps=self.eps,mod=mod)
        for j in range(ut.W):
            for i in range(h,ut.H):
                if abs(ut.matrix[i][j])>ut.eps:
                    if determinant or inverse_matrix:
                        det*=ut.matrix[i][j]
                        if mod:
                            det%=mod
                    if mod:
                        inve=MOD(mod).Pow(ut.matrix[i][j],-1)
                    else:
                        inve=1/ut.matrix[i][j]

                    ut.Elem_Raw_Operate_1(i,h)
                    if determinant and i!=h and mod:
                        det=(-det)%mod
                    if inverse_matrix:
                        im.Elem_Raw_Operate_1(i,h)
                    if linear_equation:
                        le.Elem_Raw_Operate_1(i,h)

                    ut.Elem_Raw_Operate_2(h,inve)
                    if inverse_matrix:
                        im.Elem_Raw_Operate_2(h,inve)
                    if linear_equation:
                        le.Elem_Raw_Operate_2(h,inve)

                    for ii in range(ut.H):
                        if ii==h:
                            continue
                        x=-ut.matrix[ii][j]
                        ut.Elem_Raw_Operate_3(ii,h,x)
                        if inverse_matrix:
                            im.Elem_Raw_Operate_3(ii,h,x)
                        if linear_equation:
                            le.Elem_Raw_Operate_3(ii,h,x)
                    h+=1
                    break
            else:
                det=0
        if any(le[i][0] for i in range(h,self.H)):
            le=None
        tpl=()
        if determinant:
            tpl+=(det,)
        if inverse_matrix:
            if det==0:
                im=None
            tpl+=(im,)
        if linear_equation:
            tpl+=(le,)
        if rank:
            tpl+=(h,)
        if upper_triangular:
            tpl+=(ut,)
        if len(tpl)==1:
            tpl=tpl[0]
        return tpl

mod=998244353
N,K=map(int,readline().split())
M=Matrix(2*K**2,2*K**2)
for i in range(K):
    for j in range(K):
        for k in range(K):
            if len({i,j,k})==3 and j in (min(i,j,k),max(i,j,k)):
                M[i*K+j][j*K+k]+=1
                M[K**2+i*K+j][K**2+j*K+k]+=1
                M[i*K+j][K**2+j*K+k]+=k
A=Matrix(1,2*K**2,mod=mod)
for i in range(K**2):
    A[0][i]=1
    A[0][i+K**2]+=i//K+i%K
A@=M@(N-2)
print(sum(A[0][:K**2])%mod,sum(A[0][K**2:])%mod)
yukicoder

結果

テストケース

ソースコード
