結果

問題 No.2996 Floor Sum
コンテスト
ユーザー 👑 testestest
提出日時 2024-12-15 12:27:39
言語 PyPy3
(7.3.15)
結果
TLE  
実行時間 -
コード長 1,580 bytes
コンパイル時間 370 ms
コンパイル使用メモリ 82,296 KB
実行使用メモリ 207,388 KB
最終ジャッジ日時 2024-12-20 23:30:45
合計ジャッジ時間 34,781 ms
ジャッジサーバーID
（参考情報） 		judge1 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
sample AC * 1 TLE * 1
other AC * 8 TLE * 3
権限があれば一括ダウンロードができます

ソースコード

diff # 

from fractions import Fraction as F
from functools import cache
from math import factorial as fact

def scary_sum_naive(p,q,N,M,A,B):
  return sum(i**p*((A*i+B)//M)**q for i in range(N))

def eval(f,x):
  ret=0
  for c in f[::-1]:
    ret=ret*x+c
  return ret

def newton_interpolation(x,y):
  ret=[]
  f=[1]
  for xx,yy in zip(x,y):
    e=eval(f,xx)
    c=eval(ret,xx)
    coef = (yy-c)*F(1,e)
    ret.append(0)
    for i in range(len(f)): ret[i]+=coef*f[i]
    for i in range(len(f)-1,0,-1): f[i]=f[i]*-xx+f[i-1]
    f[0]=f[0]*-xx
    f.append(1)
  while ret[-1]==0: ret.pop()
  return ret

S=[
	newton_interpolation(list(range(k+2)),[sum(i**k for i in range(N)) for N in range(k+2)])
	for k in range(23)
]

@cache
def choose(n,r,s):
  return fact(n)//fact(r)//fact(s)//fact(n-r-s)

def scary_sum(p,q,N,M,A,B):
  A1,A2=A//M,A%M
  B1,B2=B//M,B%M
  ans=0
  for r in range(q+1):
    for s in range(q-r+1):
      coef = choose(q,r,s) * A1**s * B1**(q-r-s)
      if coef!=0:
        ans+=coef*scary_sum_reduced(p+s,r,N,M,A2,B2)
  return ans

@cache
def scary_sum_reduced(p,q,N,M,A,B):
  # base case
  if N==0: return 0
  if q==0: return int(eval(S[p],N))
  K=(A*(N-1)+B)//M
  if K==0: return 0  # K==0 かつ q==0 のときは0でない
  ans = K**q * int(eval(S[p],N))
  for i in range(q):
    for j in range(p+2):
      coef=choose(q,i,0)*S[p][j]
      if coef!=0:
        ans-=coef*scary_sum(i,j,K,A,M,M-B+A-1)
  return int(ans)

T=int(input())
for _ in range(T):
	temp=list(map(int,input().split()))
	temp[2]+=1
	print(scary_sum(*temp)%998244353)
	scary_sum_reduced.cache_clear()
0