MojoでProject Euler 42

https://projecteuler.net/problem=42

String.splitあるんですね。前からありましたかね。

import sys

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


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

fn get_words(data: String) raises -> List[String]:
    var v = data.split(",")
    for s in v:
        s[] = s[].replace('"', "")
    return v

fn to_Int(word: String) -> Int:
    var sc = 0
    for c in word.as_bytes():
        sc += int(c[] - ord("A") + 1)
    return sc

fn int_sqrt(n: Int) -> Int:
    var x = 1
    x = (x + n // x) // 2
    while True:
        var x1 = (x + n // x) // 2
        if x1 >= x:
            return x
        x = x1

fn is_triangle(n: Int) -> Bool:
    # n = x(x + 1) / 2
    # x^2 + x - 2n = 0
    # x = (-1 + sqrt(1 + 8n)) / 2
    var x = (-1 + int_sqrt(1 + 8 * n)) // 2
    return x * (x + 1) // 2 == n

fn f(path: String) -> Int:
    var counter = 0
    try:
        with open(path, "r") as f:
            var data = f.read()
            var words = get_words(data)
            for word in words:
                var value = to_Int(word[])
                if is_triangle(value):
                    counter += 1
    except:
        return 0
    
    return counter

fn main():
    print(f("0042_words.txt"))

MojoでProject Euler 41

https://projecteuler.net/problem=41

nが9でも8でも1～nのpandigitalは9で割り切れるから、n=7を考えればよいので、10⁷までエラトステネスの篩をして、上からpandigitalで素数のものを探せばよいです。
ただ、一般の進法を考えると、Millar-Rabinを使えばよいでしょう。その場合はどうなるでしょう。例えばB進法と考えると、自然数は

と表されるから、ある自然数dでBを割って1余るなら、dで割った余りが各桁の和をdで割った余りと同じになって、この和はnだけで決まるのでこの余りが0ならn桁のB進自然数は素数でないことになります（nが小さいときはちょっと怪しいですが）。
こうやって計算量を減らすことができますが、すぐにオーバーフローしてしまうのであまり大きなBでは計算できません。ちゃんと計算できるのはB=19までっぽいです。

import sys


#################### prime number ####################

fn add(a: Int, b: Int, d: Int) -> Int:
    var c = a + b
    if c < 0:   # overflow
        c -= d
    elif c >= d:
        c -= d
    return c

fn mul2(a: Int, b: Int, d: Int, e1: Int, e2: Int) -> Int:
    var n = 0
    var a1 = a
    for i in range(0, e2, e1):
        var b1 = (b >> i) & ((1 << e1) - 1)
        n += a1 * b1 % d
        a1 = (a1 << e1) % d
    return n % d

fn mul3(a: Int, b: Int, d: Int, e1: Int, e2: Int) -> Int:
    var n = 0
    var a1 = a
    for i in range(0, e2, e1):
        var b1 = (b >> i) & ((1 << e1) - 1)
        n = add(n, a1 * b1 % d, d)
        a1 = (a1 << e1) % d
    return n

fn mul4(a: Int, b: Int, d: Int) -> Int:
    var n = 0
    var a1 = a
    for i in range(0, 63):
        var b1 = (b >> i) & 1
        if b1 == 1:
            n = add(n, a1, d)
        a1 = add(a1, a1, d)
    return n

fn mul(a: Int, b: Int, d: Int) -> Int:
    if d < 1 << 31:
        return a * b % d
    elif d < 1 << 40:
        var b1 = b & ((1 << 20) - 1)
        var b2 = b >> 20
        var n1 = a * b1
        var n2 = (a * b2 % d) << 20
        return (n1 + n2) % d
    elif d < 1 << 45:
        return mul2(a, b, d, 15, 45)
    elif d < 1 << 48:
        return mul2(a, b, d, 12, 48)
    elif d < 1 << 50:
        return mul2(a, b, d, 10, 50)
    elif d < 1 << 54:
        return mul3(a, b, d, 6, 54)
    elif d < 1 << 55:
        return mul3(a, b, d, 5, 55)
    elif d < 1 << 56:
        return mul3(a, b, d, 4, 56)
    elif d < 1 << 57:
        return mul3(a, b, d, 3, 57)
    elif d < 1 << 60:
        return mul3(a, b, d, 2, 60)
    else:
        return mul4(a, b, d)

fn pow(n: Int, e: Int, d: Int) -> Int:
    if e == 0:
        return 1
    elif e == 1:
        return n
    elif (e & 1) == 1:
        return mul(n, pow(n, e-1, d), d)
    else:
        var m = pow(n, e>>1, d)
        return mul(m, m, d)

fn Miller_test(n: Int, owned d: Int, a: Int) -> Bool:
    var y = pow(a, d, n)
    if y == 1:
        return True
    d <<= 1
    while d < n:
        if y == n - 1:
            return True
        y = mul(y, y, n)
        d <<= 1
    return False

fn Miller_Rabin(n: Int) -> Bool:
        var d = (n - 1) >> 1
        while (d & 1) == 0:
            d >>= 1
        
        if not Miller_test(n, d, 2):
            return False
        if not Miller_test(n, d, 3):
            return False
        if not Miller_test(n, d, 5):
            return False
        if not Miller_test(n, d, 7):
            return False
        if not Miller_test(n, d, 11):
            return False
        if not Miller_test(n, d, 13):
            return False
        if not Miller_test(n, d, 17):
            return False
        return True

fn is_prime(n: Int) -> Bool:
    if n == 2:
        return True
    elif (n & 1) == 0 or n == 1:
        return False
    else:
        for p in range(3, 20, 2):
            if p * p > n:
                return True
            elif n % p == 0:
                return False
        
        return Miller_Rabin(n)


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

trait Printable(CollectionElement, Stringable):
    pass

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 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("[]")

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 add_list[T: CollectionElement](a: List[T], b: List[T]) -> List[T]:
    var c = List[T]()
    for e in a:
        c.append(e[])
    for e in b:
        c.append(e[])
    return c

fn extend_list[T: CollectionElement](inout a: List[T], b: List[T]):
    for e in b:
        a.append(e[])

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

fn permutations[T: CollectionElement](owned i: Int, v: List[T]) -> List[T]:
    var n = len(v)
    
    var w = copy_list(v)
    var rs = List[Int]()
    for k in range(2, n+1):
        var r = i % k
        i //= k
        rs.append(r)
    
    for k in range(0, n-1):
        var r = rs[n-k-2]
        if r == 0:
            continue
        var tmp = w[k]
        w[k] = w[k+r]
        for j in range(k+r, k+1, -1):
            w[j] = w[j-1]
        w[k+1] = tmp
    
    return w

fn number(ds: List[Int], B: Int) -> Int:
    var n = 0
    for d in ds:
        n = n * B + d[]
    return n

fn str_num(owned n: Int, B: Int) -> String:
    var s: String = ""
    while n > 0:
        var d = n % B
        if d < 10:
            s = chr(d + 48) + s
        else:
            s = chr(d + 87) + s
        n //= B
    return s


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

fn is_all_not_prime(n: Int, B: Int) -> Bool:
    var s = n * (n + 1) // 2
    for d in range(2, B):
        if (B - 1) % d == 0 and s % d == 0:
            return True
    else:
        return False

fn factorial(n: Int) -> Int:
    return 1 if n == 0 else n * factorial(n - 1)

fn f_each(n: Int, B: Int) -> Int:
    var v = List[Int]()
    for d in range(n, 0, -1):
        v.append(d)
    
    for i in range(factorial(n)):
        var ds = permutations(i, v)
        var n = number(ds, B)
        if is_prime(n):
            return n
    
    return 0

fn f(B: Int) -> Int:
    for n in range(B-1, 1, -1):
        if is_all_not_prime(n, B):
            continue
        var p = f_each(n, B)
        if p > 0:
            return p
    return 0

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

MojoでProject Euler 40

https://projecteuler.net/problem=40

nが何桁のところにあるか分かればかんたんですね。

import sys

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

fn digits(owned n: Int) -> List[Int]:
    var ds = List[Int]()
    while n > 0:
        var d = n % 10
        n //= 10
        ds.append(d)
    return ds


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

fn d(n: Int) -> Int:
    var s = 0
    for e in range(1, 19):
        var s1 = s + 9 * 10**(e-1) * e
        if n <= s1:
            var k = (n - s - 1) // e + 10**(e-1)
            var r = (n - s - 1) % e
            var ds = digits(k)
            return ds[e-r-1]
        s = s1
    return 0

fn f(N: Int) -> Int:
    var s = 1
    for e in range(N):
        s *= d(10**e)
    return s

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

MojoでProject Euler 39

https://projecteuler.net/problem=39

直角三角形の各辺は



なので、となります。なので、約数が多い方が同じpで違う(k, m, n)の組が多そうです。ただ、そんなに簡単ではありません。
まず、m+nは奇数です。mは奇数でも偶数でも可能ですが、を満たさなければなりません。またmとm+nは互いに素です。kは任意です。例えば、p = 60で考えると、奇素数は3と5なので、3と5をkとmとm+nに分配すると、が考えられます。条件を満たすのはのみですが、mには2を1回だけ掛けることができるので（2は2個あるが、の頭に一つ取られるので一つだけ掛けることができる）、も条件を満たします。
ただ、2が十分たくさんあって例えば4つあれば、が条件を満たします。はいくつ2を掛けても満たしません。3と5を分配するときにそれぞれのどこかに所属するので、3通りずつあって3 * 3 = 9通りあるのですが、なので(1, 1)以外はm < m+nになるほうだけ取るので半分になって、(3*3-1)/2 = 4通りあります。
一般に奇素数のべき乗をeとすると、mに1～e個、m+nに1～e個またはmとm+nには一つも分配しないと2e+1通りあります。なので、2が十分にあれば、とすれば個組合せがあります。
つまり、pの素因数分解の指数が分かれば(k, m, n)の組の個数は上から押さえることができるので、ある個数のpがあると分かっていれば、指数を見たらそれを超えられないことが分かるかもしれません。実際にとすると、たぶん使われる素数が小さい方が個数が多くなるので、例えばの個数を数えてみると8つあるので、例えばなら(5 * 3 - 1) / 2 = 7だから、この形の素因数分解ならもう実際に数えなくても8を上回ることができないことが分かります。
なので、各指数の組合せで1000を超えないのが可能なものを列挙して、その中で素数がなるべく小さいものについて実際に(k, m, n)の組合せを数えて、これを超える可能性がある指数の組合せだけを考えます。指数の順番を変えたものと素数を大きくしたものも同じように個数を上から押さえることができます。例えば、ならも同じですし、も同じです。
これで計算してみると、で少し時間オーバーするくらいでした。
ただ、この2つ目のような例は大きな素数も試さないといけないです。実際のところはなるべく素数が小さいほうが組合せの個数が多くなるのですが、必ずしもそうとは言えないのが難しいところです。ただ、2が少なくなると個数が少なくなることは分かります。例えば、とすると、となるので、の9通りは成り立たないことが分かって、元々13通りなので、最大でも4個の組合せしかないことが分かって、8個を超えられないから、2が1つしかないときは最初から考えなくてもよく、素数の組合せを減らすことができます。いつかできたらと思います。

from collections import Set
import sys

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

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

fn make_prime_table(N: Int) -> List[Int]:
    var a = initialize_list(N+1, True)
    for p in range(2, N+1):
        if p * p > N:
            break
        elif not a[p]:
            continue
        
        for n in range(p*p, N+1, p):
            a[n] = False
    
    var ps = List[Int]()
    for n in range(2, N+1):
        if a[n]:
            ps.append(n)
    return ps


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

trait Printable(CollectionElement, Stringable):
    pass

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 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("[]")

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 add_list[T: CollectionElement](a: List[T], b: List[T]) -> List[T]:
    var c = List[T]()
    for e in a:
        c.append(e[])
    for e in b:
        c.append(e[])
    return c

fn extend_list[T: CollectionElement](inout a: List[T], b: List[T]):
    for e in b:
        a.append(e[])

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


#################### Factors ####################

struct Factors(ComparableCollectionElement):
    var ps: List[Int]
    var es: List[Int]
    var value: Int
    
    fn __init__(inout self, ps: List[Int], es: List[Int]):
        self.ps = ps
        self.es = es
        self.value = 1
        self.value = self.calc_value()
    
    fn __copyinit__(inout self, other: Factors):
        self.ps = other.ps
        self.es = other.es
        self.value = other.value
    
    fn __moveinit__(inout self, owned other: Factors):
        self.ps = other.ps^
        self.es = other.es^
        self.value = other.value
    
    fn __len__(self) -> Int:
        return len(self.ps)
    
    fn __lt__(self, other: Self) -> Bool:
        return self.value < other.value
    
    fn __le__(self, other: Self) -> Bool:
        return self.value <= other.value
    
    fn __eq__(self, other: Self) -> Bool:
        return self.value == other.value
    
    fn __ne__(self, other: Self) -> Bool:
        return self.value != other.value
    
    fn __gt__(self, other: Self) -> Bool:
        return self.value > other.value
    
    fn __ge__(self, other: Self) -> Bool:
        return self.value >= other.value
    
    fn __mul__(self, other: Factors) -> Factors:
        var ps = List[Int]()
        var es = List[Int]()
        var L1 = self.ps.size
        var L2 = other.ps.size
        var k = 0
        var l = 0
        while k < L1 and l < L2:
            var p1 = self.ps[k]
            var e1 = self.es[k]
            var p2 = other.ps[l]
            var e2 = other.es[l]
            if p1 == p2:
                ps.append(p1)
                es.append(e1+e2)
                k += 1
                l += 1
            elif p1 < p2:
                ps.append(p1)
                es.append(e1)
                k += 1
            else:
                ps.append(p2)
                es.append(e2)
                l += 1
        
        for k1 in range(k, L1):
            ps.append(self.ps[k1])
            es.append(self.es[k1])
        for l1 in range(l, L2):
            ps.append(other.ps[l1])
            es.append(other.es[l1])
        return Factors(ps, es)
    
    fn __str__(self) -> String:
        if len(self) == 0:
            return "1"
        
        var s = str(self.ps[0])
        if self.es[0] > 1:
            s += "^" + str(self.es[0])
        for i in range(1, len(self)):
            s += " * " + str(self.ps[i])
            if self.es[i] > 1:
                s += "^" + str(self.es[i])
        return s
    
    fn calc_value(self) -> Int:
        var n = 1
        for i in range(len(self.ps)):
            var p = self.ps[i]
            var e = self.es[i]
            n *= p**e
        return n
    
    @staticmethod
    fn create(n: Int) -> Factors:
        var ps = List[Int]()
        var es = List[Int]()
        var m = n
        for p in range(2, n+1):
            if p*p > m:
                break
            elif m%p == 0:
                var a = div_pow(m, p)
                var e = a.get[0, Int]()
                m = a.get[1, Int]()
                ps.append(p)
                es.append(e)
        if m > 1:
            ps.append(m)
            es.append(1)
        return Factors(ps, es)


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

fn divide_into_two(fs: Factors) -> Tuple[List[Int], List[Int]]:
    if len(fs) == 1:
        var p = fs.ps[0]
        var e = fs.es[0]
        var ds1 = List[Int](1)
        var ds2 = List[Int](1)
        for e1 in range(1, e+1):
            ds1.append(1)
            ds2.append(p**e1)
            ds1.append(p**e1)
            ds2.append(1)
        return (ds1, ds2)
    else:
        var N = len(fs)
        var mid = N // 2
        var fs1 = Factors(sublist(fs.ps, 0, mid), sublist(fs.es, 0, mid))
        var fs2 = Factors(sublist(fs.ps, mid, N), sublist(fs.es, mid, N))
        var pair1 = divide_into_two(fs1)
        var ds11 = pair1.get[0, List[Int]]()
        var ds12 = pair1.get[1, List[Int]]()
        var pair2 = divide_into_two(fs2)
        var ds21 = pair2.get[0, List[Int]]()
        var ds22 = pair2.get[1, List[Int]]()
        var ds1 = List[Int]()
        var ds2 = List[Int]()
        for i in range(len(ds11)):
            for j in range(len(ds21)):
                ds1.append(ds11[i] * ds21[j])
                ds2.append(ds12[i] * ds22[j])
        return (ds1, ds2)

fn num_right_angle_triangles(fs: Factors) -> Int:
    var e2 = fs.es[0]
    var n2 = 1 << e2
    
    # 2を削除
    var ps = List[Int]()
    var es = List[Int]()
    for i in range(1, len(fs)):
        ps.append(fs.ps[i])
        es.append(fs.es[i])
    var fs1 = Factors(ps, es)
    
    var ds = divide_into_two(fs1)
    var ds1 = ds.get[0, List[Int]]()
    var ds2 = ds.get[1, List[Int]]()
    var counter = 0
    for i in range(len(ds1)):
        var d1 = ds1[i]
        var d2 = ds2[i]
        if d1 < d2 and d1 * n2 > d2:
            counter += 1
    return counter

fn next_factors(fs: Factors, primes: List[Int]) -> List[Factors]:
    var fss = List[Factors]()
    
    # 左端は常に積める
    var e3 = fs.es[0]
    var ps1 = copy_list(fs.ps)
    var es1 = copy_list(fs.es)
    es1[0] += 1
    fss.append(Factors(ps1, es1))
    
    # 左端と違う指数が前と一つ差なら積める
    for i in range(1, len(fs)):
        var e = fs.es[i]
        if e == e3 - 1:
            var ps2 = copy_list(fs.ps)
            var es2 = copy_list(fs.es)
            es2[i] += 1
            fss.append(Factors(ps2, es2))
            break
        elif e != e3:
            break
    else:
        # 全て1なら右に伸ばせる
        if e3 == 1:
            var ps3 = copy_list(fs.ps)
            var es3 = copy_list(fs.es)
            ps3.append(primes[len(fs)+1])
            es3.append(1)
            fss.append(Factors(ps3, es3))
    return fss

fn create_shapes(N: Int, primes: List[Int]) -> List[List[Int]]:
    var ess = List[List[Int]]()
    var stack = List[Factors]()
    stack.append(Factors.create(3))
    while len(stack) > 0:
        var fs = stack.pop()
        ess.append(fs.es)
        for fs1 in next_factors(fs, primes):
            if fs1[].value * 2 <= N:
                stack.append(fs1[])
    return ess

# Nを超えないように2を掛ける
fn mul_two_pow(fs: Factors, N: Int) -> Factors:
    var ps = List[Int](2)
    var es = List[Int](0)
    extend_list(ps, fs.ps)
    extend_list(es, fs.es)
    var value = fs.value
    while value * 2 <= N:
        es[0] += 1
        value *= 2
    return Factors(ps, es)

# esの形で左から詰まった素因数分解を作る
# es: [2, 1], N: 1000 => 2^4*3^2*5
fn create_minimal_factors(es1: List[Int], N: Int, primes: List[Int]) -> Factors:
    var ps = List[Int]()
    var es = List[Int]()
    for i in range(len(es1)):
        if es1[i] != 0:
            ps.append(primes[i+1])  # 2は飛ばす
            es.append(es1[i])
    var fs = Factors(ps, es)
    return mul_two_pow(fs, N)

fn normalize_es(fs: Factors, primes: List[Int]) -> List[Int]:
    var es = List[Int]()
    # 2は飛ばす
    var i = 1
    var j = 1
    var max_p = fs.ps[len(fs)-1]
    while primes[i] <= max_p:
        if primes[i] == fs.ps[j]:
            es.append(fs.es[j])
            j += 1
        else:
            es.append(0)
        i += 1
    return es

fn swap_es(fs: Factors, i: Int) -> Factors:
    var ps = copy_list(fs.ps)
    var es = copy_list(fs.es)
    var tmp = es[i]
    es[i] = es[i+1]
    es[i+1] = tmp
    return Factors(ps, es)

fn possible_max_num_right_triangles(es: List[Int]) -> Int:
    var n = 1
    for e in es:
        n *= e[]*2 + 1
    return n // 2

fn collect_all_permutations(fs0: Factors, N: Int) -> List[Factors]:
    var fss = List[Factors]()
    var s = Set[Int](fs0.value)
    var stack = List[Factors](fs0)
    while len(stack) > 0:
        var fs = stack.pop()
        fss.append(fs)
        for i in range(len(fs)-1):
            if fs.es[i] > fs.es[i+1]:
                var fs1 = swap_es(fs, i)
                if fs1.value <= N and fs1.value not in s:
                    stack.append(fs1)
                    s.add(fs1.value)
    return fss

fn move_states(es: List[Int], N: Int, primes: List[Int]) -> List[List[Int]]:
    # 指数を移動できる
    var ess = List[List[Int]]()
    
    # 最後の0の手前を右に一つ移動できる
    var found_zero = False
    for i in range(len(es)-1, -1, -1):
        if es[i] == 0:
            found_zero = True
        elif found_zero:
            var es1 = copy_list(es)
            es1[i+1] = es1[i]
            es1[i] = 0
            ess.append(es1)
            break
    
    # 最後が0以外でその前が0の0回以上の連続でその前は全て0以外のときのみ
    # 最後を一つ右に進められる
    var mode = 0
    for i in range(len(es)-1, -1, -1):
        if mode == 0:
            mode = 1
        elif mode == 1:
            if es[i] != 0:
                mode = 2
        elif mode == 2:
            if es[i] == 0:
                mode = 3
                break
    
    if mode == 2:
        var es2 = copy_list(es)
        es2.append(es[len(es)-1])
        es2[len(es)-1] = 0
        ess.append(es2)
    
    return ess

fn make_factors(es0: List[Int], primes: List[Int]) -> Factors:
    var ps = List[Int]()
    var es = List[Int]()
    for i in range(len(es0)):
        if es0[i] != 0:
            ps.append(primes[i])
            es.append(es0[i])
    return Factors(ps, es)

fn max_num_right_triangles(es0: List[Int], N: Int,
                                    primes: List[Int]) -> Tuple[Int, Int]:
    var upper_num = possible_max_num_right_triangles(es0)
    var fs0 = make_factors(es0, primes)
    var max_num = num_right_angle_triangles(fs0)
    var max_p = fs0.value
    for fs in collect_all_permutations(fs0, N):
        var stack = List[List[Int]](fs[].es)
        while len(stack) > 0:
            var es = stack.pop()
            var fs = create_minimal_factors(es, N, primes)
            if fs.es[0] == 0:   # 2が無い
                continue
            if fs.value > N:
                continue
            var num = num_right_angle_triangles(fs)
            if num > max_num:
                max_num = num
                max_p = fs.value
            var ess = move_states(es, N, primes)
            for es1 in ess:
                stack.append(es1[])
    return (max_num, max_p)

fn f(N: Int) -> Int:
    var M = 1000
    var primes = make_prime_table(M)
    var ess = create_shapes(N, primes)
    var max_num_triangles = 0
    var max_p = 0
    for es in ess:
        var fs = create_minimal_factors(es[], N, primes)
        var n = num_right_angle_triangles(fs)
        if n > max_num_triangles:
            max_p = fs.value
            max_num_triangles = n
    
    # esを絞る
    var ess1 = List[List[Int]]()
    for es in ess:
        if possible_max_num_right_triangles(es[]) > max_num_triangles:
            ess1.append(es[])
    
    for es in ess1:
        var max_num = possible_max_num_right_triangles(es[])
        if max_num <= max_num_triangles:
            continue
        var fs = make_factors(es[], primes)
        var pair = max_num_right_triangles(es[], N, primes)
        var num = pair.get[0, Int]()
        var p = pair.get[1, Int]()
        if num > max_num_triangles:
            max_num_triangles = num
            max_p = p
    
    print(max_num_triangles)
    return max_p

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

AtCoder Beginner Contest 361 F

https://atcoder.jp/contests/abc361/tasks/abc361_f

これは昔からよく見る問題です。
として、2乗で表せる整数は、3乗で表せる整数はなどとなります。しかし、例えばなどと重複するのでそのまま数えてはいけないです。そこで底をべき乗で表されないものに限定します。64ならとします。そうすると、2乗で表せる整数はに限定され、3乗で表せる整数はとなります。一般にを1を除いたNまでのべき乗数とすると、

と書けます。
しかし、Pythonなら簡単なのですが、Rustだと簡単にオーバーフローするので累乗根の計算が難しいんです。なので、ならがになるので、そこはまとめて計算します。そうすると累乗根の計算が楽になります。

// x = a^b
#![allow(non_snake_case)]

use std::collections::HashMap;


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

// floor(x^(1/e))
fn root(x: i64, e: u32) -> i64 {
    let next_y = |y: i64| {
        (x / y.pow(e-1) + (e as i64 - 1) * y) / (e as i64)
    };
    
    let mut y = (x as f64).powf(1.0 / (e as f64)) as i64;
    if y.pow(e) < x {
        y = next_y(y)
    }
    loop {
        let y1 = next_y(y);
        if y <= y1 {
            break
        }
        y = y1
    }
    y
}

fn floor_log(n: i64, b: i64) -> u32 {
    if n < b {
        0
    }
    else {
        1 + floor_log(n / b, b)
    }
}


//////////////////// process ////////////////////

fn f(n: i64, memo: &mut HashMap<i64, i64>) -> i64 {
    if n < 4 {
        return 0
    }
    else if let Some(m) = memo.get(&n) {
        return *m
    }
    
    let mut s: i64 = 0;
    let mut e_mid: u32 = 0;
    for b in 2.. {
        let e1 = floor_log(n, b);
        if e1 <= 2 {
            e_mid = e1;
            break
        }
        
        let e2 = floor_log(n, b+1);
        s += ((e1 - e2) as i64) * (b - 1 - f(b, memo));
        if e1 == e2 + 1 {
            e_mid = e2;
            break
        }
    }
    
    for e in 2..e_mid+1 {
        let r = root(n, e);
        s += r - 1 - f(r, memo)
    }
    
    memo.insert(n, s);
    s
}

fn F(N: i64) -> i64 {
    let mut memo: HashMap<i64, i64> = HashMap::new();
    f(N, &mut memo) + 1
}

fn main() {
    let N: i64 = read();
    println!("{}", F(N))
}

AtCoder Beginner Contest 361 E

https://atcoder.jp/contests/abc361/tasks/abc361_e

端点から端点に移動するとき、横道は2回エッジを通りますがそうでない経路は1回しか通らないので、木の直径を求めればよいです。

// Tree and Hamilton Path 2
#![allow(non_snake_case)]


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// Graph ////////////////////

type Node = usize;
type Dist = i64;
type Edge = (Node, Node, Dist);

fn read_edge() -> Edge {
    let v: Vec<usize> = read_vec();
    let (A, B) = (v[0]-1, v[1]-1);
    let C = v[2] as Dist;
    (A, B, C)
}

type Graph = Vec<Vec<(Node, Dist)>>;

fn create_graph(N: usize, edges: Vec<Edge>) -> Graph {
    let mut graph: Graph = vec![vec![]; N];
    for (A, B, C) in edges.into_iter() {
        graph[A].push((B, C));
        graph[B].push((A, C))
    }
    graph
}

fn fathest_point(v0: Node, graph: &Graph) -> (Node, Dist) {
    let N = graph.len();
    let mut visited: Vec<bool> = vec![false; N];
    let mut dists: Vec<Dist> = vec![0; N];
    visited[v0] = true;
    let mut stack: Vec<Node> = vec![v0];
    while let Some(v) = stack.pop() {
        for &(v1, d1) in graph[v].iter() {
            if !visited[v1] {
                visited[v1] = true;
                dists[v1] = dists[v] + d1;
                stack.push(v1)
            }
        }
    }
    let (dmax, vmax) = (0..N).map(|i| (dists[i], i)).max().unwrap();
    (vmax, dmax)
}

fn diameter(graph: &Graph) -> Dist {
    let (v, _) = fathest_point(0, graph);
    let (_, d) = fathest_point(v, graph);
    d
}


//////////////////// process ////////////////////

fn read_input() -> Vec<Edge> {
    let N: usize = read();
    let edges: Vec<Edge> = (0..N-1).map(|_| read_edge()).collect();
    edges
}

fn F(edges: Vec<Edge>) -> Dist {
    let all_edge_dist = edges.iter().map(|&(_, _, d)| d).sum::<Dist>();
    let N = edges.len() + 1;
    let graph = create_graph(N, edges);
    let d = diameter(&graph);
    all_edge_dist * 2 - d
}

fn main() {
    let edges = read_input();
    println!("{}", F(edges))
}

AtCoder Beginner Contest 361 D

https://atcoder.jp/contests/abc361/tasks/abc361_d

各マスはWかBか空なので2ビットで表されます。あとはしらみつぶしで。

// Go Stone Puzzle
#![allow(non_snake_case)]

use std::collections::{VecDeque, HashSet};


//////////////////// library ////////////////////

fn read<T: std::str::FromStr>() -> T {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).ok();
    line.trim().parse().ok().unwrap()
}

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    read::<String>().split_whitespace()
            .map(|e| e.parse().ok().unwrap()).collect()
}


//////////////////// process ////////////////////

fn read_input() -> (usize, String, String) {
    let N: usize = read();
    let S: String = read();
    let T: String = read();
    (N, S, T)
}

type State = u32;

fn str_to_quat(S: String) -> State {
    let mut n: u32 = 15;    // 右端の2マス
    for c in S.chars().rev() {
        n = n * 4 + (if c == 'W' { 0 } else { 1 })
    }
    n
}

fn spaces(s: State, N: usize) -> (usize, usize) {
    let mut i1 = 100;
    let mut i2 = 0;
    for i in 0..N+2 {
        if ((s >> ((i as usize)*2)) & 3) == 3 {
            if i1 == 100 {
                i1 = i
            }
            else {
                i2 = i
            }
        }
    }
    (i1, i2)
}

fn get_cell(i: usize, s: State) -> u32 {
    (s >> ((i as u32) * 2)) & 3
}

fn set_cell(i: usize, s1: u32, s: State) -> State {
    let j = (i as u32) * 2;
    s ^ (((s >> j) & 3) << j) | (s1 << j)
}

fn nexts(s0: State, N: usize) -> Vec<State> {
    let mut states: Vec<State> = vec![];
    let (i1, i2) = spaces(s0, N);
    for i in 0..N+1 {
        if i == i1 || i == i2 || i+1 == i1 || i+1 == i2 {
            continue
        }
        let mut s = s0;
        let s1 = get_cell(i, s0);
        let s2 = get_cell(i+1, s0);
        s = set_cell(i1, s1, s);
        s = set_cell(i2, s2, s);
        s = set_cell(i, 3, s);
        s = set_cell(i+1, 3, s);
        states.push(s)
    }
    states
}

fn F(N: usize, S: String, T: String) -> i32 {
    if S == T {
        return 0
    }
    
    let start = str_to_quat(S);
    let end = str_to_quat(T);
    let mut visited: HashSet<State> = HashSet::new();
    let mut queue: VecDeque<(i32, State)> = VecDeque::new();
    queue.push_back((0, start));
    visited.insert(start);
    while let Some((n, s)) = queue.pop_front() {
        let ss = nexts(s, N);
        for s1 in ss.into_iter() {
            if s1 == end {
                return n+1
            }
            else if !visited.contains(&s1) {
                queue.push_back((n+1, s1));
                visited.insert(s1);
            }
        }
    }
    -1
}

fn main() {
    let (N, S, T) = read_input();
    println!("{}", F(N, S, T))
}