MojoでProject Euler 9

https://projecteuler.net/problem=9

直角三角形の各辺の長さは、

と表せて、周囲の長さLは、

となるので、L/2を3つの整数の積にすれば簡単にl, m, nを求められます。ただ、一般的には2つの約数に分解して、片方をさらに2つの約数に分解しますが、そうすると分解したときに素因数分解の形を素因数分解の形に分解しないとさらに素因数分解することになって都合が悪いです。そうすると、素因数分解のリストを作ることになりますが、それがどうしてもうまくいかないです。DynamicVectorはどうも固定サイズの要素しか受け付けないようです。
よく考えたら、最初から3つに分解すれば整数の3つ組のDynamicVectorにすればよいので、問題なくなります。

ここで、構造体に__copyinit__メソッドを定義すると代入でコピーが行われます。

struct S:
    var x: Int
    
    fn __init__(inout self, x: Int):
        self.x = x
    
    fn __copyinit__(inout self, other: Self):
        self.x = other.x

fn main():
    var s1 = S(1)
    let s2 = s1     # copy
    s1.x = 2
    print(s1.x)     # 2
    print(s2.x)     # 1

__copyinit__メソッドが無いとエラーになります。

# e009.mojo
import sys


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

fn div_pow(owned n: Int, d: Int) -> Tuple[Int, Int]:
    var e = 0
    while n % d == 0:
        n //= d
        e += 1
    return (e, n)


#################### PrimeFactors ####################

struct PrimeFactors:
    var pf: DynamicVector[Tuple[Int, Int]]
    
    fn __init__(inout self, pf: DynamicVector[Tuple[Int, Int]]):
        self.pf = pf
    
    fn __copyinit__(inout self, other: Self):
        self.pf = DynamicVector[Tuple[Int, Int]]()
        for i in range(other.size()):
            self.pf.push_back(other.pf[i])
    
    fn size(self) -> Int:
        return self.pf.size
    
    fn get(self, i: Int) -> Tuple[Int, Int]:
        return self.pf[i]
    
    fn get_p(self, i: Int) -> Int:
        return self.pf[i].get[0, Int]()
    
    fn get_e(self, i: Int) -> Int:
        return self.pf[i].get[1, Int]()
    
    fn set_e(inout self, i: Int, e1: Int):
        let p = self.get_p(i)
        let e = self.get_e(i)
        self.pf[i] = (p, e+e1)
    
    fn append(self, f: Tuple[Int, Int]) -> PrimeFactors:
        let p = f.get[0, Int]()
        let e = f.get[1, Int]()
        var pf = self
        pf.pf.push_back(f)
        return pf
    
    fn remove(inout self, i: Int) -> PrimeFactors:
        var pf = DynamicVector[Tuple[Int, Int]]()
        for j in range(i):
            pf.push_back(self.get(j))
        for j in range(i+1, self.size()):
            pf.push_back(self.get(j))
        return PrimeFactors(pf)
    
    fn find(self, p: Int) -> Int:
        for i in range(self.size()):
            if self.get_p(i) == p:
                return i
        return -1
    
    fn add(inout self, f: Tuple[Int, Int]) -> PrimeFactors:
        let p = f.get[0, Int]()
        let e = f.get[1, Int]()
        let i = self.find(p)
        if i == -1:
            return self.append(f)
        
        let e1 = self.get_e(i)
        let new_e = e1 + e
        if new_e == 0:
            return self.remove(i)
        else:
            var pf = self
            pf.set_e(i, new_e)
            return pf
    
    fn __mul__(self, other: Self) -> Self:
        var pf = self
        for i in range(other.size()):
            pf = pf.add(other.get(i))
        return pf
    
    fn __div__(self, other: Self) -> Self:
        var pf = self
        for i in range(other.size()):
            let f = other.get(i)
            let p = f.get[0, Int]()
            let e = f.get[1, Int]()
            pf = pf.add((p, -e))
        return pf
    
    fn value(self) -> Int:
        var n: Int = 1
        for i in range(self.size()):
            n *= self.get_p(i) ** self.get_e(i)
        return n
    
    fn left(self, mid: Int) -> PrimeFactors:
        var pf = DynamicVector[Tuple[Int, Int]]()
        for i in range(mid):
            pf.push_back(self.pf[i])
        return PrimeFactors(pf)
    
    fn right(self, mid: Int) -> PrimeFactors:
        var pf = DynamicVector[Tuple[Int, Int]]()
        for i in range(mid, self.size()):
            pf.push_back(self.pf[i])
        return PrimeFactors(pf)
    
    @staticmethod
    fn create() -> PrimeFactors:
        let pf = DynamicVector[Tuple[Int, Int]]()
        return PrimeFactors(pf)
    
    @staticmethod
    fn create_item(p: Int, e: Int) -> PrimeFactors:
        if e == 0:
            return PrimeFactors.one()
        else:
            var pf = PrimeFactors.create()
            pf.pf.push_back((p, e))
            return pf
    
    @staticmethod
    fn one() -> PrimeFactors:
        return PrimeFactors.create()
    
    @staticmethod
    fn factorize(n: Int) -> PrimeFactors:
        if n == 1:
            return PrimeFactors.one()
        
        for p in range(2, n+1):
            if p*p > n:
                break
            elif n % p == 0:
                let ret = div_pow(n, p)
                let e = ret.get[0, Int]()
                let m = ret.get[1, Int]()
                return PrimeFactors.create_item(p, e) * PrimeFactors.factorize(m)
        return PrimeFactors.create_item(n, 1)

alias Triplet = Tuple[Int, Int, Int]

# pf -> [(l, m, n)] ((m, n) = 1)
fn collect_tripartitions(pf : PrimeFactors) -> DynamicVector[Triplet]:
    var ds = DynamicVector[Triplet]()
    if pf.size() == 0:
        ds.push_back((1, 1, 1))
        return ds
    elif pf.size() == 1:
        let p = pf.get_p(0)
        let e = pf.get_e(0)
        for e1 in range(e+1):
            ds.push_back((p**e1, p**(e-e1), 1))
            if e != e1:
                ds.push_back((p**e1, 1, p**(e-e1)))
    else:
        let mid = pf.size() // 2
        let divs1 = collect_tripartitions(pf.left(mid))
        let divs2 = collect_tripartitions(pf.right(mid))
        for i in range(divs1.size):
            let d10 = divs1[i].get[0, Int]()
            let d11 = divs1[i].get[1, Int]()
            let d12 = divs1[i].get[2, Int]()
            for j in range(divs2.size):
                let d20 = divs2[j].get[0, Int]()
                let d21 = divs2[j].get[1, Int]()
                let d22 = divs2[j].get[2, Int]()
                ds.push_back((d10*d20, d11*d21, d12*d22))
    return ds

fn collect_triplets(L: Int) -> DynamicVector[Triplet]:
    var triplets = DynamicVector[Triplet]()
    let pf = PrimeFactors.factorize(L//2)
    let divs = collect_tripartitions(pf)
    for i in range(len(divs)):
        let l = divs[i].get[0, Int]()
        let m = divs[i].get[1, Int]()
        let p = divs[i].get[2, Int]()
        let n = p - m
        if n <= 0 or m <= n:
            continue
        triplets.push_back((l*(m*m-n*n), 2*l*m*n, l*(m*m+n*n)))
    return triplets

fn f(L: Int):
    if L % 2 == 1:
        return
    
    let triplets = collect_triplets(L)
    for i in range(len(triplets)):
        let a = triplets[i].get[0, Int]()
        let b = triplets[i].get[1, Int]()
        let c = triplets[i].get[2, Int]()
        print(a * b * c)

fn main() raises:
    let args = sys.argv()
    let L = atol(args[1])
    f(L)