結果

問題 No.215 素数サイコロと合成数サイコロ (3-Hard)
コンテスト
ユーザー Min_25
提出日時 2018-08-04 03:04:49
言語 Python3
(3.13.1 + numpy 2.2.1 + scipy 1.14.1)
結果
RE  
(最新)
AC  
(最初)
実行時間 -
コード長 1,766 bytes
コンパイル時間 75 ms
コンパイル使用メモリ 12,800 KB
実行使用メモリ 44,352 KB
最終ジャッジ日時 2024-09-19 17:42:26
合計ジャッジ時間 3,252 ms
ジャッジサーバーID
（参考情報） 		judge4 / judge5
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other RE * 2
権限があれば一括ダウンロードができます

ソースコード

diff # 

import numpy as np
from numpy.fft import fft, ifft

def poly_mul_mod(f, g, mod, shift=15):
  def conv(f):
    ff = np.zeros(fft_size, dtype=np.complex128)
    ff[:len(f)] = ((f & mask) + (f >> shift) * 1j) / 2.0
    ff = fft(ff)
    ffrc = np.concatenate((ff[0:1], ff[-1:0:-1])).conj()
    return ff, ffrc

  mask = (1 << shift) - 1
  s = len(f) + len(g) - 1
  fft_size = 1 << ((2 * s - 1).bit_length() - 1)
  ff, ffrc = conv(f)
  fg, fgrc = conv(g)
  ffr, ffi = ff + ffrc, ff - ffrc
  fgr, fgi = fg + fgrc, fg - fgrc
  f01 = ifft(ffr * (fgr + fgi) + ffi * fgr)[:s]
  lo = f01.real.round().astype(np.int64)
  mid = f01.imag.round().astype(np.int64)
  hi = ifft(-ffi * fgi).real.round().astype(np.int64)[:s]
  ret = (lo + ((mid % mod) << shift) + ((hi % mod) << (2 * shift))) % mod
  return ret

def nth(n, numer, denom, mod):
  while n > 0:
    sdenom = denom.copy()
    sdenom[1::2] = mod - denom[1::2]
    numer = poly_mul_mod(numer, sdenom, mod)[n & 1::2]
    denom = poly_mul_mod(denom, sdenom, mod)[::2]
    n >>= 1
  return numer[0]

def solve():
  import sys

  Ps = np.array([2, 3, 5, 7, 11, 13], dtype=np.int)
  Cs = np.array([4, 6, 8, 9, 10, 12], dtype=np.int)
  mod = 10 ** 9 + 7

  def gene(ds, T):
    dp = np.zeros((T + 1, ds[-1] * T + 1), dtype=np.int)
    dp[0][0] = 1
    o = ds[0]
    for di in range(6):
      d = ds[di]
      for t in range(T):
        dp[t+1][d+o*t:d*(t+1)+1] = \
          (dp[t+1][d+o*t:d*(t+1)+1] + dp[t][o*t:d*t+1]) % mod
    return dp[T][:]

  for line in sys.stdin:
    N, P, C = map(int, line.split())
    denom = poly_mul_mod(gene(Ps, P), gene(Cs, C), mod)
    denom = (mod - denom) % mod
    denom[0] = 1
    numer = np.cumsum(denom, dtype=np.int64) % mod
    print(nth(N + len(denom) - 1, numer, denom, mod))

solve()
0