#708397 (Python3) No.718 行列のできるフィボナッチ数列道場 (1)

提出ソース
結果

問題	No.718 行列のできるフィボナッチ数列道場 (1)
ユーザー	vwxyz
提出日時	2021-10-20 23:47:32
言語	Python3 (3.12.2 + numpy 1.26.4 + scipy 1.12.0)
結果	AC
実行時間	48 ms / 2,000 ms
コード長	17,644 bytes
コンパイル時間	248 ms
コンパイル使用メモリ	14,464 KB
実行使用メモリ	13,440 KB
最終ジャッジ日時	2024-09-20 06:48:02
合計ジャッジ時間	2,137 ms
ジャッジサーバーID （参考情報）	judge3 / judge2
このコードへのチャレンジ
（要ログイン）
テストケース

テストケース表示
入力	結果	実行時間実行使用メモリ
testcase_00	AC	47 ms 13,440 KB
testcase_01	AC	47 ms 13,440 KB
testcase_02	AC	47 ms 13,312 KB
testcase_03	AC	46 ms 13,312 KB
testcase_04	AC	46 ms 13,312 KB
testcase_05	AC	47 ms 13,440 KB
testcase_06	AC	46 ms 13,440 KB
testcase_07	AC	46 ms 13,440 KB
testcase_08	AC	48 ms 13,440 KB
testcase_09	AC	44 ms 13,440 KB
testcase_10	AC	44 ms 13,440 KB
testcase_11	AC	44 ms 13,312 KB
testcase_12	AC	44 ms 13,312 KB
testcase_13	AC	43 ms 13,312 KB
testcase_14	AC	43 ms 13,312 KB
testcase_15	AC	47 ms 13,312 KB
testcase_16	AC	47 ms 13,312 KB
testcase_17	AC	46 ms 13,312 KB
testcase_18	AC	46 ms 13,440 KB
testcase_19	AC	44 ms 13,440 KB
testcase_20	AC	43 ms 13,312 KB
testcase_21	AC	46 ms 13,312 KB
testcase_22	AC	46 ms 13,312 KB
権限があれば一括ダウンロードができます
ソースコード

raw source code
import bisect
import copy
import decimal
import fractions
import functools
import heapq
import itertools
import math
import random
import sys
from collections import Counter,deque,defaultdict
from functools import lru_cache,reduce
from heapq import heappush,heappop,heapify,heappushpop,_heappop_max,_heapify_max
def _heappush_max(heap,item):
    heap.append(item)
    heapq._siftdown_max(heap, 0, len(heap)-1)
def _heappushpop_max(heap, item):
    if heap and item < heap[0]:
        item, heap[0] = heap[0], item
        heapq._siftup_max(heap, 0)
    return item
from math import gcd as GCD
read=sys.stdin.read
readline=sys.stdin.readline
readlines=sys.stdin.readlines

class Polynomial:
    def __init__(self,polynomial,max_degree=-1,eps=0,mod=0):
        self.max_degree=max_degree
        if self.max_degree!=-1 and len(polynomial)>self.max_degree+1:
            self.polynomial=polynomial[:self.max_degree+1]
        else:
            self.polynomial=polynomial
        self.mod=mod
        self.eps=eps

    def __eq__(self,other):
        if type(other)!=Polynomial:
            return False
        if len(self.polynomial)!=len(other.polynomial):
            return False
        for i in range(len(self.polynomial)):
            if self.eps<abs(self.polynomial[i]-other.polynomial[i]):
                return False
        return True

    def __ne__(self,other):
        if type(other)!=Polynomial:
            return True
        if len(self.polynomial)!=len(other.polynomial):
            return True
        for i in range(len(self.polynomial)):
            if self.eps<abs(self.polynomial[i]-other.polynomial[i]):
                return True
        return False

    def __add__(self,other):
        if type(other)==Polynomial:
            summ=[0]*max(len(self.polynomial),len(other.polynomial))
            for i in range(len(self.polynomial)):
                summ[i]+=self.polynomial[i]
            for i in range(len(other.polynomial)):
                summ[i]+=other.polynomial[i]
            if self.mod:
                for i in range(len(summ)):
                    summ[i]%=self.mod
        else:
            summ=[x for x in self.polynomial] if self.polynomial else [0]
            summ[0]+=other
            if self.mod:
                summ[0]%=self.mod
        while summ and abs(summ[-1])<=self.eps:
            summ.pop()
        summ=Polynomial(summ,max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        return summ

    def __sub__(self,other):
        if type(other)==Polynomial:
            diff=[0]*max(len(self.polynomial),len(other.polynomial))
            for i in range(len(self.polynomial)):
                diff[i]+=self.polynomial[i]
            for i in range(len(other.polynomial)):
                diff[i]-=other.polynomial[i]
            if self.mod:
                for i in range(len(diff)):
                    diff[i]%=self.mod
        else:
            diff=[x for x in self.polynomial] if self.polynomial else [0]
            diff[0]-=other
            if self.mod:
                diff[0]%=self.mod
        while diff and abs(diff[-1])<=self.eps:
            diff.pop()
        diff=Polynomial(diff,max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        return diff

    def __mul__(self,other):
        if type(other)==Polynomial:
            if self.max_degree==-1:
                prod=[0]*(len(self.polynomial)+len(other.polynomial)-1)
                for i in range(len(self.polynomial)):
                    for j in range(len(other.polynomial)):
                        prod[i+j]+=self.polynomial[i]*other.polynomial[j]
            else:
                prod=[0]*min(len(self.polynomial)+len(other.polynomial)-1,self.max_degree+1)
                for i in range(len(self.polynomial)):
                    for j in range(min(len(other.polynomial),self.max_degree+1-i)):
                        prod[i+j]+=self.polynomial[i]*other.polynomial[j]
            if self.mod:
                for i in range(len(prod)):
                    prod[i]%=self.mod
        else:
            if self.mod:
                prod=[x*other%self.mod for x in self.polynomial]
            else:
                prod=[x*other for x in self.polynomial]
        while prod and abs(prod[-1])<=self.eps:
            prod.pop()
        prod=Polynomial(prod,max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        return prod

    def __matmul__(self,other):
        assert type(other)==Polynomial
        if self.mod:
            prod=NTT(self.polynomial,other.polynomial)
        else:
            prod=FFT(self.polynomial,other.polynomial)
        if self.max_degree!=-1 and len(prod)>self.max_degree+1:
            prod=prod[:self.max_degree+1]
            while prod and abs(prod[-1])<=self.eps:
                prod.pop()
        prod=Polynomial(prod,max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        return prod

    def __truediv__(self,other):
        if type(other)==Polynomial:
            assert other.polynomial
            for n in range(len(other.polynomial)):
                if self.eps<abs(other.polynomial[n]):
                    break
            assert len(self.polynomial)>n
            for i in range(n):
                assert abs(self.polynomial[i])<=self.eps
            self_polynomial=self.polynomial[n:]
            other_polynomial=other.polynomial[n:]
            if self.mod:
                inve=MOD(self.mod).Pow(other_polynomial[0],-1)
            else:
                inve=1/other_polynomial[0]
            quot=[]
            for i in range(len(self_polynomial)-len(other_polynomial)+1):
                if self.mod:
                    quot.append(self_polynomial[i]*inve%self.mod)
                else:
                    quot.append(self_polynomial[i]*inve)
                for j in range(len(other_polynomial)):
                    self_polynomial[i+j]-=other_polynomial[j]*quot[-1]
                    if self.mod:
                        self_polynomial[i+j]%=self.mod
            for i in range(max(0,len(self_polynomial)-len(other_polynomial)+1),len(self_polynomial)):
                if self.eps<abs(self_polynomial[i]):
                    assert self.max_degree!=-1
                    self_polynomial=self_polynomial[-len(other_polynomial)+1:]+[0]*(len(other_polynomial)-1-len(self_polynomial))
                    while len(quot)<=self.max_degree:
                        self_polynomial.append(0)
                        if self.mod:
                            quot.append(self_polynomial[0]*inve%self.mod)
                            self_polynomial=[(self_polynomial[i]-other_polynomial[i]*quot[-1])%self.mod for i in range(1,len(self_polynomial))]
                        else:
                            quot.append(self_polynomial[0]*inve)
                            self_polynomial=[(self_polynomial[i]-other_polynomial[i]*quot[-1]) for i in range(1,len(self_polynomial))]
                    break
            quot=Polynomial(quot,max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        else:
            assert self.eps<abs(other)
            if self.mod:
                inve=MOD(self.mod).Pow(other,-1)
                quot=Polynomial([x*inve%self.mod for x in self.polynomial],max_degree=self.max_degree,eps=self.eps,mod=self.mod)
            else:
                quot=Polynomial([x/other for x in self.polynomial],max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        return quot

    def __rtruediv__(self,other):
        assert self.polynomial and self.eps<self.polynomial[0]
        assert self.max_degree!=-1
        if self.mod:
            quot=[MOD(self.mod).Pow(self.polynomial[0],-1)]
            if self.mod==998244353:
                prim_root=3
                prim_root_inve=332748118
            else:
                prim_root=Primitive_Root(self.mod)
                prim_root_inve=MOD(self.mod).Pow(prim_root,-1)
            def DFT(polynomial,n,inverse=False):
                polynomial=polynomial+[0]*((1<<n)-len(polynomial))
                if inverse:
                    for bit in range(1,n+1):
                        a=1<<bit-1
                        x=pow(prim_root,mod-1>>bit,mod)
                        U=[1]
                        for _ in range(a):
                            U.append(U[-1]*x%mod)
                        for i in range(1<<n-bit):
                            for j in range(a):
                                s=i*2*a+j
                                t=s+a
                                polynomial[s],polynomial[t]=(polynomial[s]+polynomial[t]*U[j])%mod,(polynomial[s]-polynomial[t]*U[j])%mod
                    x=pow((mod+1)//2,n,mod)
                    for i in range(1<<n):
                        polynomial[i]*=x
                        polynomial[i]%=mod
                else:
                    for bit in range(n,0,-1):
                        a=1<<bit-1
                        x=pow(prim_root_inve,mod-1>>bit,mod)
                        U=[1]
                        for _ in range(a):
                            U.append(U[-1]*x%mod)
                        for i in range(1<<n-bit):
                            for j in range(a):
                                s=i*2*a+j
                                t=s+a
                                polynomial[s],polynomial[t]=(polynomial[s]+polynomial[t])%mod,U[j]*(polynomial[s]-polynomial[t])%mod
                return polynomial
        else:
            quot=[1/self.polynomial[0]]
            def DFT(polynomial,n,inverse=False):
                N=len(polynomial)
                if inverse:
                    primitive_root=[math.cos(-i*2*math.pi/(1<<n))+math.sin(-i*2*math.pi/(1<<n))*1j for i in range(1<<n)]
                else:
                    primitive_root=[math.cos(i*2*math.pi/(1<<n))+math.sin(i*2*math.pi/(1<<n))*1j for i in range(1<<n)]
                polynomial=polynomial+[0]*((1<<n)-N)
                if inverse:
                    for bit in range(1,n+1):
                        a=1<<bit-1
                        for i in range(1<<n-bit):
                            for j in range(a):
                                s=i*2*a+j
                                t=s+a
                                polynomial[s],polynomial[t]=polynomial[s]+polynomial[t]*primitive_root[j<<n-bit],polynomial[s]-polynomial[t]*primitive_root[j<<n-bit]
                    for i in range(1<<n):
                        polynomial[i]=round((polynomial[i]/(1<<n)).real)
                else:
                    for bit in range(n,0,-1):
                        a=1<<bit-1
                        for i in range(1<<n-bit):
                            for j in range(a):
                                s=i*2*a+j
                                t=s+a
                                polynomial[s],polynomial[t]=polynomial[s]+polynomial[t],primitive_root[j<<n-bit]*(polynomial[s]-polynomial[t])

                return polynomial
        for n in range(self.max_degree.bit_length()):
            prev=quot
            if self.mod:
                polynomial=[x*y*y%self.mod for x,y in zip(DFT(self.polynomial[:1<<n+1],n+2),DFT(prev,n+2))]
                quot=DFT(polynomial,n+2,inverse=True)[:1<<n+1]
            else:
                polynomial=[x*y*y for x,y in zip(DFT(self.polynomial[:1<<n+1],n+2),DFT(prev,n+2))]
                quot=DFT(polynomial,n+2,inverse=True)[:1<<n+1]
            for i in range(1<<n):
                quot[i]=2*prev[i]-quot[i]
                if self.mod:
                    quot[i]%=self.mod
            for i in range(1<<n,1<<n+1):
                quot[i]=-quot[i]
                if self.mod:
                    quot[i]%=self.mod
        quot=quot[:self.max_degree+1]
        for i in range(len(quot)):
            quot[i]*=other
            if self.mod:
                quot[i]%=self.mod
        return quot

    def __floordiv__(self,other):
        assert type(other)==Polynomial
        quot=[0]*(len(self.polynomial)-len(other.polynomial)+1)
        rema=[x for x in self.polynomial]
        if self.mod:
            inve=MOD(self.mod).Pow(other.polynomial[-1],-1)
            for i in range(len(self.polynomial)-len(other.polynomial),-1,-1):
                quot[i]=rema[i+len(other.polynomial)-1]*inve%self.mod
                for j in range(len(other.polynomial)):
                    rema[i+j]-=quot[i]*other.polynomial[j]
                    rema[i+j]%=self.mod
        else:
            inve=1/other.polynomial[-1]
            for i in range(len(self.polynomial)-len(other.polynomial),-1,-1):
                quot[i]=rema[i+len(other.polynomial)-1]*inve
                for j in range(len(other.polynomial)):
                    rema[i+j]-=quot[i]*other.polynomial[j]
        quot=Polynomial(quot,max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        return quot

    def __mod__(self,other):
        assert type(other)==Polynomial
        quot=[0]*(len(self.polynomial)-len(other.polynomial)+1)
        rema=[x for x in self.polynomial]
        if self.mod:
            inve=MOD(self.mod).Pow(other.polynomial[-1],-1)
            for i in range(len(self.polynomial)-len(other.polynomial),-1,-1):
                quot[i]=rema[i+len(other.polynomial)-1]*inve%self.mod
                for j in range(len(other.polynomial)):
                    rema[i+j]-=quot[i]*other.polynomial[j]
                    rema[i+j]%=self.mod
        else:
            inve=1/other.polynomial[-1]
            for i in range(len(self.polynomial)-len(other.polynomial),-1,-1):
                quot[i]=rema[i+len(other.polynomial)-1]*inve
                for j in range(len(other.polynomial)):
                    rema[i+j]-=quot[i]*other.polynomial[j]
        while rema and abs(rema[-1])<=self.eps:
            rema.pop()
        rema=Polynomial(rema,max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        return rema

    def __divmod__(self,other):
        assert type(other)==Polynomial
        quot=[0]*(len(self.polynomial)-len(other.polynomial)+1)
        rema=[x for x in self.polynomial]
        if self.mod:
            inve=MOD(self.mod).Pow(other.polynomial[-1],-1)
            for i in range(len(self.polynomial)-len(other.polynomial),-1,-1):
                quot[i]=rema[i+len(other.polynomial)-1]*inve%self.mod
                for j in range(len(other.polynomial)):
                    rema[i+j]-=quot[i]*other.polynomial[j]
                    rema[i+j]%=self.mod
        else:
            inve=1/other.polynomial[-1]
            for i in range(len(self.polynomial)-len(other.polynomial),-1,-1):
                quot[i]=rema[i+len(other.polynomial)-1]*inve
                for j in range(len(other.polynomial)):
                    rema[i+j]-=quot[i]*other.polynomial[j]
        while rema and abs(rema[-1])<=self.eps:
            rema.pop()
        quot=Polynomial(quot,max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        rema=Polynomial(rema,max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        return quot,rema

    def __neg__(self):
        if self.mod:
            nega=Polynomial([(-x)%self.mod for x in self.polynomial],max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        else:
            nega=Polynomial([-x for x in self.polynomial],max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        return nega

    def __pos__(self):
        posi=Polynomial([x for x in self.polynomial],max_degree=self.max_degree,eps=self.eps,mod=self.mod)
        return posi

    def __bool__(self):
        return self.polynomial

    def __getitem__(self,n):
        if n<=len(self.polynomial)-1:
            return self.polynomial[n]
        else:
            return 0

    def __setitem__(self,n,x):
        if self.mod:
            x%=self.mod
        if self.max_degree==-1 or n<=self.max_degree:
            if n<=len(self.polynomial)-1:
                self.polynomial[n]=x
            elif self.eps<abs(x):
                self.polynomial+=[0]*(n-len(self.polynomial))+[x]

    def __call__(self,x):
        retu=0
        pow_x=1
        for i in range(len(self.polynomial)):
            retu+=pow_x*self.polynomial[i]
            pow_x*=x
            if self.mod:
                retu%=self.mod
                pow_x%=self.mod
        return retu

    def __str__(self):
        return "["+", ".join(map(str,self.polynomial))+"]"

    def degree(self):
        return len(self.polynomial)-1

def Bostan_Mori(poly_nume,poly_deno,N,mod=0,fft=False,ntt=False):
    if type(poly_nume)==Polynomial:
        poly_nume=poly_nume.polynomial
    if type(poly_deno)==Polynomial:
        poly_deno=poly_deno.polynomial
    if ntt:
        convolve=NTT
    elif fft:
        convolve=FFT
    else:
        def convolve(poly_nume,poly_deno):
            conv=[0]*(len(poly_nume)+len(poly_deno)-1)
            for i in range(len(poly_nume)):
                for j in range(len(poly_deno)):
                    conv[i+j]+=poly_nume[i]*poly_deno[j]
            if mod:
                for i in range(len(conv)):
                    conv[i]%=mod
            return conv
    while N:
        poly_deno_=[-x if i%2 else x for i,x in enumerate(poly_deno)]
        if N%2:
            poly_nume=convolve(poly_nume,poly_deno_)[1::2]
        else:
            poly_nume=convolve(poly_nume,poly_deno_)[::2]
        poly_deno=convolve(poly_deno,poly_deno_)[::2]
        if fft and mod:
            for i in range(len(poly_nume)):
                poly_nume[i]%=mod
            for i in range(len(poly_deno)):
                poly_deno[i]%=mod
        N//=2
    return poly_nume[0]

N=int(readline())
mod=10**9+7
ans=Bostan_Mori([0,1],[1,-1,-1],N,mod=mod)*Bostan_Mori([0,1],[1,-1,-1],N+1,mod=mod)%mod
print(ans)
yukicoder

結果

テストケース

ソースコード
