MojoでProject Euler 65

https://projecteuler.net/problem=65

連分数の係数fに対して、連分数を途中で打ち切って通常の分数にしたとき、例えば、

 \displaystyle x = f_0 + \frac{1}{f_1 + \frac{1}{f_2}}

としたとき、その手前で打ち切ったときは \frac{f_0f_1+1}{f_1}、さらにその手前で打ち切ったときは f_0となりますが、それぞれの分子分母を a_1,\ b_1,\ a_0,\ b_0とすると、

 a_2 = f_2a_1 + a_0
 b_2 = f_2b_1 + b_0

と計算できます。これは一般的に言えます。
オーバーフローするので、多倍長整数を使わなければなりません。

import sys


#################### List ####################

fn initialize_list[T: CollectionElement](N: Int, init: T) -> List[T]:
    var a = List[T](capacity=N)
    for n in range(N):
        a.append(init)
    return a

trait Printable(CollectionElement, Stringable):
    pass

fn print_list[T: Printable](a: List[T]):
    if a.size > 0:
        var s = "[" + str(a[0])
        for i in range(1, a.size):
            s += ", " + str(a[i])
        s += "]"
        print(s)
    else:
        print("[]")


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

@value
struct BigInteger(Stringable):
    var v: List[Int]
    
    fn __eq__(self, other: BigInteger) -> Bool:
        if len(self.v) != len(other.v):
            return False
        for i in range(len(self.v)):
            if self.v[i] != other.v[i]:
                return False
        return True
    
    fn __lt__(self, other: BigInteger) -> Bool:
        if len(self.v) != len(other.v):
            return len(self.v) < len(other.v)
        for i in range(len(self.v)-1, -1, -1):
            if self.v[i] != other.v[i]:
                return self.v[i] < other.v[i]
        return False
    
    fn __gt__(self, other: BigInteger) -> Bool:
        if len(self.v) != len(other.v):
            return len(self.v) > len(other.v)
        for i in range(len(self.v)-1, -1, -1):
            if self.v[i] != other.v[i]:
                return self.v[i] > other.v[i]
        return False
    
    fn __ge__(self, other: BigInteger) -> Bool:
        if len(self.v) != len(other.v):
            return len(self.v) > len(other.v)
        for i in range(len(self.v)-1, -1, -1):
            if self.v[i] != other.v[i]:
                return self.v[i] > other.v[i]
        return True
    
    fn __add__(self, other: BigInteger) -> BigInteger:
        var v = List[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.append(n % 10)
            carry = n // 10
        if carry > 0:
            v.append(carry)
        return BigInteger(v)
    
    # 非負になる前提
    fn __sub__(self, other: BigInteger) -> BigInteger:
        var v = List[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.append(n % 10)
            carry = n // 10
        if v.size > 1 and v[v.size-1] == 0:
            var tmp = v.pop()   # 受けないとwarning
        
        return BigInteger(v)
    
    fn __mul__(self, other: Int) -> BigInteger:
        var v = List[Int]()
        var carry = 0
        for d in self.v:
            var n = d[] * other + carry
            v.append(n % 10)
            carry = n // 10
        while carry > 0:
            var r = carry % 10
            carry //= 10
            v.append(r)
        return BigInteger(v)
    
    fn __str__(self) -> String:
        if len(self.v) == 0:
            return "0"
        
        var s: String = ""
        for i in range(self.v.size-1, -1, -1):
            if self.v[i] < 10:
                s += chr(self.v[i] + 48)
            else:
                s += chr(self.v[i] + 87)
        return s
    
    @staticmethod
    fn create(owned n: Int) -> BigInteger:
        var v = List[Int]()
        while n > 0:
            var d = n % 10
            n //= 10
            v.append(d)
        return BigInteger(v)


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

fn f(N: Int) -> Int:
    var a = initialize_list(N, BigInteger.create(2))
    a[1] = BigInteger.create(3)
    for i in range(2, N):
        var f = 1 if i % 3 != 2 else (i // 3 + 1) * 2
        a[i] = a[i-1] * f + a[i-2]
    
    var s = 0
    var c = a[N-1].v
    for d in c:
        s += d[]
    return s

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