MojoでProject Euler 58

https://projecteuler.net/problem=58

右上が 4n^2-2n+1、左上が 4n^2+1、左下が 4n^2+2n+1なので、 n^2+1と同じようなやり方で篩が使えます。
0.1だと一瞬ですが、この割合が小さくなると急速に時間がかかるようになり、0.06でも4秒くらいでした。

import sys


#################### process ####################

# 4n^2 - 2n + 1
fn sieve1(N: Int) -> List[Int]:
    var ns = List[Int]()
    var a = List[Int]()
    for n in range(N+1):
        a.append(4*n*n-2*n+1)
    
    for n in range(1, N+1):
        if a[n] == 1:
            continue
        elif a[n] == 4*n*n-2*n+1:   # prime?
            ns.append(n)
        var p = a[n]
        for n1 in range(n+p, N+1, p):
            while a[n1] % p == 0:
                a[n1] //= p
        var m = ((p + 1) // 2 - n) % p
        for n2 in range(m, N+1, p):
            while a[n2] % p == 0:
                a[n2] //= p
    return ns

fn sieve2(N: Int) -> List[Int]:
    var ns = List[Int]()
    var a = List[Int]()
    for n in range(N+1):
        a.append(4*n*n+1)
    
    for n in range(1, N+1):
        if a[n] == 1:
            continue
        elif a[n] == 4*n*n+1:   # prime?
            ns.append(n)
        var p = a[n]
        for n1 in range(n+p, N+1, p):
            while a[n1] % p == 0:
                a[n1] //= p
        var m = (p - n) % p
        for n2 in range(m, N+1, p):
            while a[n2] % p == 0:
                a[n2] //= p
    return ns

# 4n^2 + 2n + 1
fn sieve3(N: Int) -> List[Int]:
    var ns = List[Int]()
    var a = List[Int]()
    for n in range(N+1):
        a.append(4*n*n+2*n+1)
    
    for n in range(1, N+1):
        if a[n] == 1:
            continue
        elif a[n] == 4*n*n+2*n+1:   # prime?
            ns.append(n)
        var p = a[n]
        for n1 in range(n+p, N+1, p):
            while a[n1] % p == 0:
                a[n1] //= p
        var m = ((p - 1) // 2 - n) % p
        for n2 in range(m, N+1, p):
            while a[n2] % p == 0:
                a[n2] //= p
    return ns

fn g(ratio: Float64, N: Int) -> Int:
    var ns1 = sieve1(N)
    var ns2 = sieve2(N)
    var ns3 = sieve3(N)
    var num_primes = 0
    var num_diagonals = 1
    var k1 = 0
    var k2 = 0
    var k3 = 0
    for i in range(1, N+1):
        if k1 < len(ns1) and ns1[k1] == i:
            num_primes += 1
            k1 += 1
        if k2 < len(ns2) and ns2[k2] == i:
            num_primes += 1
            k2 += 1
        if k3 < len(ns3) and ns3[k3] == i:
            num_primes += 1
            k3 += 1
        num_diagonals += 4
        if num_primes / num_diagonals < ratio:
            return i * 2 + 1
    else:
        return 0

fn f(ratio: Float64) -> Int:
    var N = 100000
    while True:
        var n = g(ratio, N)
        if n > 0:
            return n
        N *= 2

fn main() raises:
    var args = sys.argv()
    var ratio = atof(args[1])
    print(f(ratio))