結果

問題 No.732 3PrimeCounting
ユーザー tktk_snsntktk_snsn
提出日時 2020-12-17 23:21:10
言語 Python3
(3.13.1 + numpy 2.2.1 + scipy 1.14.1)
結果
TLE  
(最新)
AC  
(最初)
実行時間 -
コード長 1,078 bytes
コンパイル時間 162 ms
コンパイル使用メモリ 12,800 KB
実行使用メモリ 112,964 KB
最終ジャッジ日時 2024-09-21 08:37:14
合計ジャッジ時間 8,931 ms
ジャッジサーバーID
（参考情報） 		judge2 / judge3
このコードへのチャレンジ
（要ログイン）
ファイルパターン 結果
other TLE * 1 -- * 88
権限があれば一括ダウンロードができます

ソースコード

diff # 

from itertools import chain
from scipy.signal import fftconvolve
import numpy as np


def convolve(A, B):
    sz = len(A) + len(B)
    fft_len = 1 << (sz-1).bit_length()
    F1 = np.fft.rfft(A, fft_len)
    F2 = np.fft.rfft(B, fft_len)
    G = F1 * F2
    res = np.fft.irfft(G, fft_len)
    return np.rint(res + 0.001).astype(np.int64)[:sz - 1]


def prime_set(N):
    """
    Nまでの素数のsetを返す
    """
    if N < 4:
        return ({}, {}, {2}, {2, 3})[N]
    Nsq = int(N ** 0.5 + 0.5) + 1
    primes = {2, 3} | set(chain(range(5, N + 1, 6), range(7, N + 1, 6)))
    for i in range(5, Nsq, 2):
        if i in primes:
            primes -= set(range(i * i, N + 1, i * 2))
    return primes


N = int(input())
primes = np.array(list(prime_set(N)), dtype=np.int32)
F = np.bincount(1 * primes).astype(np.int64)
G = np.bincount(2 * primes).astype(np.int64)
H = np.bincount(3 * primes).astype(np.int64)

FFF = convolve(F, convolve(F, F))
FG = convolve(F, G)

cnt = (FFF - 3 * FG + 2 * H) // 6
p = np.array(list(prime_set(len(cnt))), dtype=np.int32)
print(cnt[p].sum())
0