MojoでProject Euler 80

https://projecteuler.net/problem=80

2ならにして、平方根を取って小数点以下を切り捨てます。

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

fn copy_list[T: CollectionElement](a: List[T]) -> List[T]:
    return sublist(a, 0, len(a))

fn sublist[T: CollectionElement](a: List[T], first: Int, last: Int) -> List[T]:
    var b = List[T]()
    for i in range(first, last):
        b.append(a[i])
    return b

fn reverse_list[T: CollectionElement](a: List[T]) -> List[T]:
    var b = List[T]()
    for i in range(len(a) - 1, -1, -1):
        b.append(a[i])
    return b

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 copy(self) -> BigInteger:
        var v = copy_list(self.v)
        return BigInteger(v)
    
    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 __ne__(self, other: BigInteger) -> Bool:
        return not self.__eq__(other)
    
    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
        
        while len(v) > 1 and v[len(v)-1] == 0:
            _ = v.pop(len(v)-1)
        
        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 __rmul__(self, other: Int) -> BigInteger:
        return self * other
    
    fn __mul__(self, other: BigInteger) -> BigInteger:
        var L1 = len(self.v)
        var L2 = len(other.v)
        var v = initialize_list(L1 + L2 - 1, 0)
        for i in range(L1):
            for j in range(L2):
                v[i+j] += self.v[i] * other.v[j]
        
        var carry = 0
        for i in range(L1 + L2 - 1):
            v[i] += carry
            carry = v[i] // 10
            v[i] = v[i] % 10
        while carry > 0:
            v.append(carry % 10)
            carry //= 10
        return BigInteger(v)
    
    fn less_than(self, other: BigInteger, offs: Int) -> Bool:
        if len(self.v) != len(other.v) + offs:
            return len(self.v) < len(other.v) + offs
        for i in range(len(other.v)-1, -1, -1):
            if self.v[i+offs] != other.v[i]:
                return self.v[i+offs] < other.v[i]
        return False
    
    fn div(self, other: BigInteger, offs: Int) -> Int:
        for i in range(1, 10):
            if self.less_than(other * i, offs):
                return i - 1
        else:
            return 9
    
    fn shift(self, offs: Int) -> BigInteger:
        var v = List[Int]()
        for _ in range(offs):
            v.append(0)
        for d in self.v:
            v.append(d[])
        return BigInteger(v)
    
    fn __floordiv__(self, other: BigInteger) -> BigInteger:
        var v = List[Int]()
        var dividend = self.copy()
        var offs = len(self.v) - len(other.v)
        if self.less_than(other, offs):
            offs -= 1
        for i in range(offs, -1, -1):
            var d = dividend.div(other, i)
            dividend = dividend - (other * d).shift(i)
            v.append(d)
        return BigInteger(reverse_list(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 digits(owned n: Int) -> List[Int]:
    var ds = List[Int]()
    while n > 0:
        var d = n % 10
        ds.append(d)
        n //= 10
    return ds

fn sqrt(n: BigInteger) -> BigInteger:
    var one = BigInteger.create(1)
    var two = BigInteger.create(2)
    var x = (n + one) // two
    while True:
        var x1 = (x + n // x) // two
        if x1 >= x:
            return x
        x = x1

fn f(N: Int, L: Int) -> Int:
    var s = 0
    var r = 1
    for n in range(1, N+1):
        if r * r == n:
            r += 1
            continue
        var big_n = BigInteger.create(n).shift(L * 2 - 2)
        var big_r = sqrt(big_n)
        for d in big_r.v:
            s += d[]
        print(n, s)
    return s

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