MojoでProject Euler 30

https://projecteuler.net/problem=30

2から順に等式が成り立つかどうか調べていっても十分間に合うのですが、これは重複組み合わせを使うと計算量が減ります。1634を分解して4乗和を取ると1634になりますが、1346や6431も1634になります。4桁で重複が無いから24個の数がが同じになることが分かります。6桁だと重複組み合わせは _10H_6 = 5005通りしかありません。なので、重複組合せをだして、[1, 3, 4, 6]のとき、4乗和が1634になって、それを数字に分解してソートすると、[1, 3, 4, 6]で元と同じになるので、1634が該当する数になります。この方法で17乗まで計算できました。

最新のバージョンで、DynamicVectorはListになって、Pythonと同じような便利なメソッドも追加されました。メソッドではないですが、ソートもできるんですね。

https://docs.modular.com/mojo/stdlib/collections/list#list

from algorithm.sort import sort
from math import min
import sys


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

fn print_vector(v: List[Int]):
    for k in range((v.size + 9) // 10):
        var s = str("")
        for i in range(k*10, min(k*10+10, v.size)):
            s += str(v[i]) + " "
        print(s)

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


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

fn duplicate_combinations(a: List[Int], n: Int) -> List[List[Int]]:
    var v = List[List[Int]]()
    if n == 1:
        for e in a:
            var w = List[Int](e[])
            v.append(w)
    elif n % 2 == 1:
        var v1 = duplicate_combinations(a, n-1)
        for e in a:
            for w1 in v1:
                if e[] > w1[][0]:
                    continue
                var w = List[Int](e[])
                w.extend(w1[])
                v.append(w)
    else:
        var m = n // 2
        var v1 = duplicate_combinations(a, m)
        for w1 in v1:
            for w2 in v1:
                if w1[][m-1] > w2[][0]:
                    continue
                var w = w1[]
                w.extend(w2[])
                v.append(w)
    return v

fn upper_num_digits(E: Int) -> Int:
    for n in range(1, 20):
        if n * 9**E < 10**(n-1):
            return n - 1
    return 0

fn is_coincident(w: List[Int], v: List[Int]) -> Bool:
    var delta = v.size - w.size
    for i in range(delta):
        if v[i] != 0:
            return False
    for i in range(w.size):
        if w[i] != v[i+delta]:
            return False
    return True

fn valid_sum_pows(v: List[Int], E: Int) -> Int:
    var s = 0
    for e in v:
        s += e[]**E
    
    var w = digits(s)
    sort(w)
    
    if is_coincident(w, v):
        return s
    else:
        return 0

fn f(E: Int) -> Int:
    var N = upper_num_digits(E)
    var ds = List[Int]()
    for d in range(10):
        ds.append(d)
    var v = duplicate_combinations(ds, N)
    var s = 0
    for w in v:
        s += valid_sum_pows(w[], E)
    return s - 1

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

AtCoder Beginner Contest 350 E

https://atcoder.jp/contests/abc350/tasks/abc350_e

2番目を選んだ時、nのときの期待値は E_n = (E_n+E_{\lfloor{n/2}\rfloor} + \cdots + E_{\lfloor{n/6}\rfloor})/6 + Yなので、
\displaystyle E_n = \frac{1}{5}\sum_{b=2}^6{E_{\lfloor{n/b}\rfloor}} + \frac{6}{5}Y
となります。取り得るnのは、 \lfloor{N/a}\rfloorとしたときのaが \frac{1}{6}\log_2{N}\log_3{N}\log_5{N}程度で、これが約10000です。また、aが違っても \lfloor{N/a}\rfloorが同じことが多いので、さらに取り得るnは少なくなって、十分に速いです。
結局、
\displaystyle E_0 = 0
\displaystyle E_n = \min{(E_{\lfloor{n/2}\rfloor} + X, \frac{1}{5}\sum_{b=2}^6{E_{\lfloor{n/b}\rfloor}} + \frac{6}{5}Y)}
となります。

// Toward 0
#![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()
}


//////////////////// Game ////////////////////

use std::collections::HashMap;

struct Game {
    N: u64,
    A: u64,
    X: f64,
    Y: f64,
    E: HashMap<u64, f64>
}

impl Game {
    fn exp_value(&mut self, n: u64) -> f64 {
        if n == 0 {
            return 0.0
        }
        
        match self.E.get(&n) {
            Some(e) => *e,
            None    => {
                let e1 = self.exp_value(n/self.A) + self.X;
                let e2 = (2..7).map(|b| self.exp_value(n/b)).sum::<f64>() * 0.2 +
                                                                self.Y * 1.2;
                let e = if e1 < e2 { e1 } else { e2 };
                self.E.insert(n, e);
                e
            }
        }
    }
    
    fn read() -> Game {
        let v = read_vec();
        let (N, A, X, Y) = (v[0], v[1], v[2] as f64, v[3] as f64);
        let E = HashMap::new();
        Game { N, A, X, Y, E }
    }
}

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

fn F(mut game: Game) -> f64 {
    game.exp_value(game.N)
}

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

AtCoder Beginner Contest 350 D

https://atcoder.jp/contests/abc350/tasks/abc350_d

連結成分に分解して、個々のグラフが完全グラフにするのに何本エッジが要るかを調べます。

// New Friends
#![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 ////////////////////

use std::collections::HashMap;

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

fn read_edge() -> Edge {
    let v = read_vec::<Node>();
    (v[0]-1, v[1]-1)
}

struct Graph {
    g: Vec<Vec<Node>>
}

impl Graph {
    fn len(&self) -> usize {
        self.g.len()
    }
    
    fn num_edges(&self) -> usize {
        self.g.iter().map(|vs| vs.len()).sum::<usize>() / 2
    }
    
    fn divide_into_connected(&self) -> Vec<Graph> {
        let N = self.g.len();
        let mut u = UnionFind::new(N);
        for (n, v) in self.g.iter().enumerate() {
            for &m in v.iter() {
                u.join(n, m)
            }
        }
        
        let mut m: HashMap<Node, Vec<Node>> = HashMap::new();
        for v in 0..N {
            let r = u.root(v);
            let e = m.entry(r).or_insert(vec![]);
            (*e).push(v)
        }
        
        m.into_iter().
            map(|(_, vs)| Graph { g: vs.into_iter().map(|v| self.g[v].to_vec()).
                                                collect::<Vec<Vec<Node>>>() }).
            collect::<Vec<Graph>>()
    }
    
    fn create(N: usize, edges: Vec<Edge>) -> Graph {
        let mut g: Vec<Vec<Node>> = (0..N).map(|_| vec![]).collect();
        for (u, v) in edges.into_iter() {
            g[u].push(v);
            g[v].push(u)
        }
        Graph { g }
    }
}


//////////////////// UnionFind ////////////////////

use std::cmp::max;

struct UnionFind {
    parents: Vec<Node>,
    heights: Vec<i32>,
}

impl UnionFind {
    fn new(N: usize) -> UnionFind {
        let parents: Vec<Node> = (0..N).collect();
        let heights: Vec<i32> = vec![1; N];
        UnionFind { parents, heights }
    }
    
    fn is_root(&self, v: Node) -> bool {
        self.parents[v] == v
    }
    
    fn join(&mut self, u: Node, v: Node) {
        let r1 = self.root(u);
        let r2 = self.root(v);
        if r1 == r2 {
            return  // ここには来ないはず
        }
        
        let h1 = self.heights[r1];
        let h2 = self.heights[r2];
        if h1 <= h2 {   // r2にr1がぶら下がる
            self.parents[r1] = r2;
            self.heights[r2] = max(self.heights[r2], self.heights[r1]+1);
        }
        else {
            self.parents[r2] = r1;
            self.heights[r1] = max(self.heights[r1], self.heights[r2]+1);
        }
    }
    
    fn root(&self, v0: Node) -> Node {
        if self.is_root(v0) {
            v0
        }
        else {
            let p = self.parents[v0];
            self.root(p)
        }
    }
}


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

fn read_input() -> (usize, Vec<Edge>) {
    let v = read_vec();
    let (N, M) = (v[0], v[1]);
    let edges: Vec<Edge> = (0..M).map(|_| read_edge()).collect();
    (N, edges)
}

fn F(N: usize, edges: Vec<Edge>) -> usize {
    let graph = Graph::create(N, edges);
    let subgraphs = graph.divide_into_connected();
    subgraphs.iter().map(|g| (g.len() * (g.len()-1) / 2 - g.num_edges())).
                                                            sum::<usize>()
}

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

MojoでProject Euler 29（3）

https://projecteuler.net/problem=29

前回は例えば2乗までのとき何個がダブるかをナイーブに数えていましたが、2乗までなら2～N/2がダブると分かるので、数えるまでもありません。6乗までだと、N/6～N/3の間は2から5の倍数はダブりますが、重複を考えると包除原理を使わないといけません。しかし、前回よりかなり速いはずです。最後のところで多倍長整数を使うと 10^9より大きいときも計算できます。 10^{15}でも30秒程度でした。実際のところ、 \lfloor\log_2{N}\rfloorが素数でないと速いです。

from collections import Dict
from math import min, max, abs
import sys


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

fn gcd(n: Int, m: Int) -> Int:
    return n if m == 0 else gcd(m, n % m)

fn lcm(n: Int, m: Int) -> Int:
    return n // abs(gcd(n, m)) * m

# floor(n^(1/e))
fn int_root(n: Int, e: Int) -> Int:
    if e == 1:
        return n
    
    # 3より小さければ2
    var m = 1
    for e1 in range(e):
        m *= 3
        if m > n:
            return 2
    
    var x = n
    while True:
        var x1 = n
        for i in range(e-1):
            x1 //= x
        x1 += (e-1)*x
        x1 //= e
        if x1 >= x:
            break
        x = x1
    
    while x**e > n:
        x -= 1
    while x**e <= n:
        x += 1
    return x - 1

fn int_log(b: Int, n: Int) -> Int:
    for e in range(1, n+1):
        if b**e > n:
            return e - 1
    else:
        return n

fn initialize_vector(N: Int, init: Int) -> DynamicVector[Int]:
    var a = DynamicVector[Int]()
    for n in range(N):
        a.push_back(init)
    return a

fn print_vector(v: DynamicVector[Int]):
    var s = String("[")
    if v.size > 0:
        s = s + str(v[0])
    for i in range(1, v.size):
        s = s + ", " + str(v[i])
    s = s + "]"
    print(s)


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

struct BigInteger(CollectionElement):
    var v: DynamicVector[Int]
    
    fn __init__(inout self, v: DynamicVector[Int]):
        self.v = v
    
    fn __copyinit__(inout self, other: BigInteger):
        self.v = other.v
    
    fn __moveinit__(inout self, owned other: BigInteger):
        self.v = other.v^
    
    fn __add__(self, other: BigInteger) -> BigInteger:
        var v = DynamicVector[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.push_back(n % 10)
            carry = n // 10
        if carry > 0:
            v.push_back(carry)
        return BigInteger(v)
    
    # 非負になる前提
    fn __sub__(self, other: BigInteger) -> BigInteger:
        var v = DynamicVector[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.push_back(n % 10)
            carry = n // 10
        if v.size > 1 and v[v.size-1] == 0:
            var tmp = v.pop_back()  # 受けないとwarning
        
        return BigInteger(v)
    
    fn __mul__(self, other: Int) -> BigInteger:
        var v = DynamicVector[Int]()
        var carry = 0
        for d in self.v:
            var n = d[] * other + carry
            v.push_back(n % 10)
            carry = n // 10
        while carry > 0:
            var r = carry % 10
            carry //= 10
            v.push_back(r)
        return BigInteger(v)
    
    fn __str__(self) -> String:
        var s: String = ""
        for i in range(self.v.size-1, -1, -1):
            s += String(self.v[i])
        return s
    
    @staticmethod
    fn create(n: Int) -> BigInteger:
        var v = DynamicVector[Int]()
        v.push_back(n)
        return BigInteger(v)
    
    @staticmethod
    fn parse(s: String) -> BigInteger:
        var v = DynamicVector[Int]()
        for i in range(len(s)-1, -1, -1):
            v.push_back(ord(s[i]) - 48)
        return BigInteger(v)


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

def create_pow_freq(N: Int) -> DynamicVector[Int]:
    var max_e = int_log(2, N)
    var a = initialize_vector(max_e+1, 1)
    for e in range(max_e, 0, -1):
        var m = int_root(N, e) - 1
        a[e] = m
        for e1 in range(e*2, max_e+1, e):
            a[e] -= a[e1]
    
    var b = initialize_vector(max_e+1, 1)
    for e in range(1, max_e):
        b[e] = a[e] - a[e+1]
    
    return b

fn add_d(d: Int, ds: DynamicVector[Int]) -> DynamicVector[Int]:
    var new_ds = DynamicVector[Int]()
    for d1 in ds:
        if d1[] % d != 0:
            new_ds.push_back(d1[])
    new_ds.push_back(d)
    return new_ds

fn inex_core(ds: DynamicVector[Int], M: Int) -> DynamicVector[Int]:
    var ns = DynamicVector[Int]()
    if ds.size == 1:
        ns.push_back(1)
        ns.push_back(-ds[0])
    else:
        var mid = ds.size // 2
        var ds1 = DynamicVector[Int]()
        var ds2 = DynamicVector[Int]()
        for i in range(ds.size):
            if i < mid:
                ds1.push_back(ds[i])
            else:
                ds2.push_back(ds[i])
        var ns1 = inex_core(ds1, M)
        var ns2 = inex_core(ds2, M)
        for n1 in ns1:
            for n2 in ns2:
                var n = lcm(n1[], n2[])
                if -M <= n <= M:
                    ns.push_back(n)
    return ns

fn inex(ds: DynamicVector[Int], M: Int) -> Int:
    if ds.size == 1:
        return M // ds[0]
    else:
        var mid = ds.size // 2
        var ds1 = DynamicVector[Int]()
        var ds2 = DynamicVector[Int]()
        for i in range(ds.size):
            if i < mid:
                ds1.push_back(ds[i])
            else:
                ds2.push_back(ds[i])
        
        var ns1 = inex_core(ds1, M)
        var ns2 = inex_core(ds2, M)
        var s = 0
        for n1 in ns1:
            for n2 in ns2:
                var n = lcm(n1[], n2[])
                if n > 0:
                    s += M // n
                else:
                    s -= M // (-n)
        return M - s

fn count_duplicated_each(ds: DynamicVector[Int],
                            e1: Int, e: Int, N: Int) -> Int:
    if e1 == 1:
        return N // e - 1
    elif ds.size == 1 and ds[0] == 1:
        return N * e1 // e - N * (e1 - 1) // e
    else:
        var n1 = inex(ds, N * e1 // e)
        var n2 = inex(ds, N * (e1 - 1) // e)
        return n1 - n2

fn count_duplicated(e: Int, N: Int) -> Int:
    var c = 0
    var ds = DynamicVector[Int]()
    for e1 in range(e-1, 0, -1):
        var d = e1 // gcd(e1, e)
        ds = add_d(d, ds)
        var c1 = count_duplicated_each(ds, e1, e, N)
        c += c1
    
    return c

fn f(N: Int) raises -> BigInteger:
    var b = create_pow_freq(N)
    
    var M = BigInteger.create(N-1)
    var counter = M * (N-1)
    var c = BigInteger.create(0)
    for e in range(b.size-1, 1, -1):
        c = c + BigInteger.create(b[e])
        counter = counter - c * count_duplicated(e, N)
    return counter

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

MojoでProject Euler 29（2）

https://projecteuler.net/problem=29

こういう問題はダブりの方を数える方が速いです。
例えば、5をベースで考えると、 25^2 = 5^4などダブりがあります。 25^{50} = 5^{100}まで、49個あります。これは6でも同じです。 36^2 = 6^4などです。
このように、100までで何乗まであるかが同じならダブりの個数は同じです。6乗まであるのは2だけ、4乗まであるのは3だけ、2乗まであるのは6つなどです。

from collections import Set, KeyElement
import sys


#################### Power ####################

@value
struct Power(KeyElement):
    var b: Int
    var e: Int
    
    fn __init__(inout self, b: Int, e: Int):
        self.b = b
        self.e = e
    
    fn pow(p: Power, e: Int) -> Power:
        return Power(p.b, p.e * e)
    
    fn __hash__(self) -> Int:
        var a = DTypePointer[DType.int8].alloc(2)
        a[0] = self.b
        a[1] = self.e
        var h = hash(a, 2)
        a.free()
        return h
    
    fn __eq__(self, other: Self) -> Bool:
        return self.b == other.b and self.e == other.e
    
    fn __str__(self) -> String:
        return str(self.b) + "^" + str(self.e)


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

fn make_power_table(N: Int) -> DynamicVector[Power]:
    var v = DynamicVector[Power]()
    for n in range(N+1):
        v.push_back(Power(n, 1))
    
    for n in range(2, N+1):
        var b = v[n].b
        var e = v[n].e
        if e != 1:
            continue
        var m = b
        for e1 in range(2, N):
            m *= b
            if m > N:
                break
            v[m] = Power(b, e1)
    return v

fn f(N: Int) -> Int:
    var table = make_power_table(N)
    var s = Set[Power]()
    for n in range(2, N+1):
        var a = table[n]
        for b in range(2, N+1):
            s.add(a.pow(b))
    return len(s)

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

MojoでProject Euler 29（1）

https://projecteuler.net/problem=29

a^bなら (a, b)としてダブりを排除します。ただし、 aがべき乗数ならべき乗を外に出します。例えば、 64^2 = 2^{12}とします。
MojoでSetを使うには、KeyElementを継承します。

from collections import Set, KeyElement
import sys


#################### Power ####################

@value
struct Power(KeyElement):
    var b: Int
    var e: Int
    
    fn __init__(inout self, b: Int, e: Int):
        self.b = b
        self.e = e
    
    fn pow(p: Power, e: Int) -> Power:
        return Power(p.b, p.e * e)
    
    fn __hash__(self) -> Int:
        var a = DTypePointer[DType.int8].alloc(2)
        a[0] = self.b
        a[1] = self.e
        var h = hash(a, 2)
        a.free()
        return h
    
    fn __eq__(self, other: Self) -> Bool:
        return self.b == other.b and self.e == other.e
    
    fn __str__(self) -> String:
        return str(self.b) + "^" + str(self.e)


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

fn make_power_table(N: Int) -> DynamicVector[Power]:
    var v = DynamicVector[Power]()
    for n in range(N+1):
        v.push_back(Power(n, 1))
    
    for n in range(2, N+1):
        var b = v[n].b
        var e = v[n].e
        if e != 1:
            continue
        var m = b
        for e1 in range(2, N):
            m *= b
            if m > N:
                break
            v[m] = Power(b, e1)
    return v

fn f(N: Int) -> Int:
    var table = make_power_table(N)
    var s = Set[Power]()
    for n in range(2, N+1):
        var a = table[n]
        for b in range(2, N+1):
            s.add(a.pow(b))
    return len(s)

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

MojoでProject Euler 28

https://projecteuler.net/problem=28

角から次の角を求められればよいです。

import sys

alias Point = Tuple[Int, Int]

fn proceed(pt: Point, n: Int) -> Tuple[Point, Int]:
    var x = pt.get[0, Int]()
    var y = pt.get[1, Int]()
    if x >= 0 and y >= 0:
        return ((x+1, -y-1), n+x*2+2)
    elif x >= 0:
        return ((-x, y), n+x*2)
    elif y < 0:
        return ((x, -y), n-x*2)
    else:
        return ((-x, y), n-x*2)

fn equals(pt1: Point, pt2: Point) -> Bool:
    var x1 = pt1.get[0, Int]()
    var y1 = pt1.get[1, Int]()
    var x2 = pt2.get[0, Int]()
    var y2 = pt2.get[1, Int]()
    return x1 == x2 and y1 == y2

fn f(N: Int) -> Int:
    var M = N // 2
    var n = 1
    var pt = (0, 0)
    var s = 1
    while not equals(pt, (M, M)):
        var t = proceed(pt, n)
        pt = t.get[0, Point]()
        n = t.get[1, Int]()
        s += n
    return s

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