MojoでProject Euler 97

https://projecteuler.net/problem=97

Pythonならpow関数を使うだけですが、10桁だとすぐにオーバーフローします。Problem 48で作ったpow関数を使います。

import sys


#################### pow ####################

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)


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

# a2^b+1
fn f(a: Int, b: Int) -> Int:
    var c = pow(2, b, 10**10)
    return (a * c + 1) % 10**10

fn main():
    print(f(28433, 7830457))