MojoでProject Euler 29（3）

https://projecteuler.net/problem=29

前回は例えば2乗までのとき何個がダブるかをナイーブに数えていましたが、2乗までなら2～N/2がダブると分かるので、数えるまでもありません。6乗までだと、N/6～N/3の間は2から5の倍数はダブりますが、重複を考えると包除原理を使わないといけません。しかし、前回よりかなり速いはずです。最後のところで多倍長整数を使うとより大きいときも計算できます。でも30秒程度でした。実際のところ、が素数でないと速いです。

from collections import Dict
from math import min, max, abs
import sys


#################### library ####################

fn gcd(n: Int, m: Int) -> Int:
    return n if m == 0 else gcd(m, n % m)

fn lcm(n: Int, m: Int) -> Int:
    return n // abs(gcd(n, m)) * m

# floor(n^(1/e))
fn int_root(n: Int, e: Int) -> Int:
    if e == 1:
        return n
    
    # 3より小さければ2
    var m = 1
    for e1 in range(e):
        m *= 3
        if m > n:
            return 2
    
    var x = n
    while True:
        var x1 = n
        for i in range(e-1):
            x1 //= x
        x1 += (e-1)*x
        x1 //= e
        if x1 >= x:
            break
        x = x1
    
    while x**e > n:
        x -= 1
    while x**e <= n:
        x += 1
    return x - 1

fn int_log(b: Int, n: Int) -> Int:
    for e in range(1, n+1):
        if b**e > n:
            return e - 1
    else:
        return n

fn initialize_vector(N: Int, init: Int) -> DynamicVector[Int]:
    var a = DynamicVector[Int]()
    for n in range(N):
        a.push_back(init)
    return a

fn print_vector(v: DynamicVector[Int]):
    var s = String("[")
    if v.size > 0:
        s = s + str(v[0])
    for i in range(1, v.size):
        s = s + ", " + str(v[i])
    s = s + "]"
    print(s)


#################### BigInteger ####################

struct BigInteger(CollectionElement):
    var v: DynamicVector[Int]
    
    fn __init__(inout self, v: DynamicVector[Int]):
        self.v = v
    
    fn __copyinit__(inout self, other: BigInteger):
        self.v = other.v
    
    fn __moveinit__(inout self, owned other: BigInteger):
        self.v = other.v^
    
    fn __add__(self, other: BigInteger) -> BigInteger:
        var v = DynamicVector[Int]()
        var carry = 0
        for i in range(max(self.v.size, other.v.size)):
            var d1 = self.v[i] if i < self.v.size else 0
            var d2 = other.v[i] if i < other.v.size else 0
            var n = d1 + d2 + carry
            v.push_back(n % 10)
            carry = n // 10
        if carry > 0:
            v.push_back(carry)
        return BigInteger(v)
    
    # 非負になる前提
    fn __sub__(self, other: BigInteger) -> BigInteger:
        var v = DynamicVector[Int]()
        var carry = 0
        for i in range(max(self.v.size, other.v.size)):
            var d1 = self.v[i] if i < self.v.size else 0
            var d2 = other.v[i] if i < other.v.size else 0
            var n = d1 - d2 + carry
            v.push_back(n % 10)
            carry = n // 10
        if v.size > 1 and v[v.size-1] == 0:
            var tmp = v.pop_back()  # 受けないとwarning
        
        return BigInteger(v)
    
    fn __mul__(self, other: Int) -> BigInteger:
        var v = DynamicVector[Int]()
        var carry = 0
        for d in self.v:
            var n = d[] * other + carry
            v.push_back(n % 10)
            carry = n // 10
        while carry > 0:
            var r = carry % 10
            carry //= 10
            v.push_back(r)
        return BigInteger(v)
    
    fn __str__(self) -> String:
        var s: String = ""
        for i in range(self.v.size-1, -1, -1):
            s += String(self.v[i])
        return s
    
    @staticmethod
    fn create(n: Int) -> BigInteger:
        var v = DynamicVector[Int]()
        v.push_back(n)
        return BigInteger(v)
    
    @staticmethod
    fn parse(s: String) -> BigInteger:
        var v = DynamicVector[Int]()
        for i in range(len(s)-1, -1, -1):
            v.push_back(ord(s[i]) - 48)
        return BigInteger(v)


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

def create_pow_freq(N: Int) -> DynamicVector[Int]:
    var max_e = int_log(2, N)
    var a = initialize_vector(max_e+1, 1)
    for e in range(max_e, 0, -1):
        var m = int_root(N, e) - 1
        a[e] = m
        for e1 in range(e*2, max_e+1, e):
            a[e] -= a[e1]
    
    var b = initialize_vector(max_e+1, 1)
    for e in range(1, max_e):
        b[e] = a[e] - a[e+1]
    
    return b

fn add_d(d: Int, ds: DynamicVector[Int]) -> DynamicVector[Int]:
    var new_ds = DynamicVector[Int]()
    for d1 in ds:
        if d1[] % d != 0:
            new_ds.push_back(d1[])
    new_ds.push_back(d)
    return new_ds

fn inex_core(ds: DynamicVector[Int], M: Int) -> DynamicVector[Int]:
    var ns = DynamicVector[Int]()
    if ds.size == 1:
        ns.push_back(1)
        ns.push_back(-ds[0])
    else:
        var mid = ds.size // 2
        var ds1 = DynamicVector[Int]()
        var ds2 = DynamicVector[Int]()
        for i in range(ds.size):
            if i < mid:
                ds1.push_back(ds[i])
            else:
                ds2.push_back(ds[i])
        var ns1 = inex_core(ds1, M)
        var ns2 = inex_core(ds2, M)
        for n1 in ns1:
            for n2 in ns2:
                var n = lcm(n1[], n2[])
                if -M <= n <= M:
                    ns.push_back(n)
    return ns

fn inex(ds: DynamicVector[Int], M: Int) -> Int:
    if ds.size == 1:
        return M // ds[0]
    else:
        var mid = ds.size // 2
        var ds1 = DynamicVector[Int]()
        var ds2 = DynamicVector[Int]()
        for i in range(ds.size):
            if i < mid:
                ds1.push_back(ds[i])
            else:
                ds2.push_back(ds[i])
        
        var ns1 = inex_core(ds1, M)
        var ns2 = inex_core(ds2, M)
        var s = 0
        for n1 in ns1:
            for n2 in ns2:
                var n = lcm(n1[], n2[])
                if n > 0:
                    s += M // n
                else:
                    s -= M // (-n)
        return M - s

fn count_duplicated_each(ds: DynamicVector[Int],
                            e1: Int, e: Int, N: Int) -> Int:
    if e1 == 1:
        return N // e - 1
    elif ds.size == 1 and ds[0] == 1:
        return N * e1 // e - N * (e1 - 1) // e
    else:
        var n1 = inex(ds, N * e1 // e)
        var n2 = inex(ds, N * (e1 - 1) // e)
        return n1 - n2

fn count_duplicated(e: Int, N: Int) -> Int:
    var c = 0
    var ds = DynamicVector[Int]()
    for e1 in range(e-1, 0, -1):
        var d = e1 // gcd(e1, e)
        ds = add_d(d, ds)
        var c1 = count_duplicated_each(ds, e1, e, N)
        c += c1
    
    return c

fn f(N: Int) raises -> BigInteger:
    var b = create_pow_freq(N)
    
    var M = BigInteger.create(N-1)
    var counter = M * (N-1)
    var c = BigInteger.create(0)
    for e in range(b.size-1, 1, -1):
        c = c + BigInteger.create(b[e])
        counter = counter - c * count_duplicated(e, N)
    return counter

fn main() raises:
    var args = sys.argv()
    var N = atol(args[1])
    print(f(N).__str__())