結果

問題 No.732 3PrimeCounting
ユーザー rlangevin
提出日時 2023-07-28 12:30:35
言語 Python3
(3.13.1 + numpy 2.2.1 + scipy 1.14.1)
結果
AC  
実行時間 663 ms / 3,000 ms
コード長 948 bytes
コンパイル時間 211 ms
コンパイル使用メモリ 12,672 KB
実行使用メモリ 86,428 KB
最終ジャッジ日時 2024-10-06 02:15:29
合計ジャッジ時間 46,137 ms
ジャッジサーバーID
(参考情報)
judge3 / judge1
このコードへのチャレンジ
(要ログイン)
ファイルパターン 結果
other AC * 89
権限があれば一括ダウンロードができます

ソースコード

diff #

import numpy as np
from math import sqrt, ceil

def Sieve(n):
    lst = [True] * (n + 1)
    S = set()
    for i in range(2, ceil(sqrt(n)) + 1):
        if lst[i]:
            for j in range(2 * i, n + 1, i):
                lst[j] = False

    for i in range(3):
        lst[i] = False
    return lst

def convolve(f, g):
    fft_len = 1
    while 2 * fft_len < len(f) + len(g) - 1:
        fft_len *= 2
    fft_len *= 2
    Ff = np.fft.rfft(f, fft_len)
    Fg = np.fft.rfft(g, fft_len)
    Fh = Ff * Fg
    h = np.fft.irfft(Fh, fft_len)
    h = np.rint(h).astype(np.int64)

    return h[:len(f) + len(g) - 1]


N = int(input())
L = Sieve(3 * N + 5)
f = np.zeros(N + 1)
f2 = np.zeros(3 * N)
for i in range(3, N + 1):
    if L[i]:
        f[i] = 1
    if L[i]:
        f2[2 * i] = 1

g = convolve(convolve(f, f), f)
h = convolve(f2, f)[:len(g)]

ans = g - 3 * h
cnt = 0
for i in range(3, len(ans)):
    if L[i]:
        cnt += ans[i]

print(cnt//6)
0