MojoでProject Euler 29(2)

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 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))

AtCoder Beginner Contest 349 E

https://atcoder.jp/contests/abc349/tasks/abc349_e

状態数が 3^9 = 19683しかなくて、深さも9しかないので、ふつうにメモ化すればよいです。

// Weighted Tic-Tac-Toe
#![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()
}


//////////////////// Color ////////////////////

#[derive(PartialEq, Copy, Clone)]
enum Color {
    White = 0,
    Red = 1,
    Blue = 2
}

impl Color {
    fn from_u32(n: u32) -> Color {
        match n {
            0 => Color::White,
            1 => Color::Red,
            _ => Color::Blue
        }
    }
}


//////////////////// State ////////////////////

struct State {
    C: [[Color; 3]; 3]
}

impl State {
    fn clone(&self) -> State {
        State { C: self.C.clone() }
    }
    
    fn exists_line(&self, color: Color) -> bool {
        for i in 0..3 {
            if (0..3).all(|j| self.C[i][j] == color) {
                return true
            }
        }
        for j in 0..3 {
            if (0..3).all(|i| self.C[i][j] == color) {
                return true
            }
        }
        if (0..3).all(|i| self.C[i][i] == color) {
            true
        }
        else if (0..3).all(|i| self.C[i][2-i] == color) {
            true
        }
        else {
            false
        }
    }
    
    fn wins_red(&self) -> bool {
        self.exists_line(Color::Red)
    }
    
    fn wins_blue(&self) -> bool {
        self.exists_line(Color::Blue)
    }
    
    fn num_remained_cells(&self) -> usize {
        self.C.iter().flatten().filter(|&c| *c == Color::White).count()
    }
    
    fn is_finished(&self) -> bool {
        self.num_remained_cells() == 0
    }
    
    fn turn(&self) -> Color {
        let n = self.num_remained_cells();
        if n % 2 == 1 { Color::Red } else { Color::Blue }
    }
    
    fn encode(&self) -> u32 {
        let mut s: u32 = 0;
        for i in (0..3).rev() {
            for j in (0..3).rev() {
                s = s * 3 + (self.C[i][j] as u32)
            }
        }
        s
    }
    
    fn decode(mut s: u32) -> State {
        let mut C: [[Color; 3]; 3] = [[Color::White; 3]; 3];
        for i in 0..3 {
            for j in 0..3 {
                C[i][j] = Color::from_u32(s % 3);
                s /= 3
            }
        }
        State { C }
    }
    
    fn nexts(&self) -> Vec<u32> {
        let mut states: Vec<u32> = vec![];
        let color = self.turn();
        for i in 0..3 {
            for j in 0..3 {
                if self.C[i][j] == Color::White {
                    let mut s = self.clone();
                    s.C[i][j] = color;
                    states.push(s.encode())
                }
            }
        }
        states
    }
}


//////////////////// Board ////////////////////

struct Board {
    A: [[i64; 3]; 3],
    win_vec: Vec<i32>
}

impl Board {
    fn wins_core(&mut self, s: u32) -> bool {
        let state = State::decode(s);
        if state.wins_red() {
            true
        }
        else if state.wins_blue() {
            false
        }
        else if state.is_finished() {
            let mut sum_red: i64 = 0;
            let mut sum_blue: i64 = 0;
            for i in 0..3 {
                for j in 0..3 {
                    if state.C[i][j] == Color::Red {
                        sum_red += self.A[i][j]
                    }
                    else {
                        sum_blue += self.A[i][j]
                    }
                }
            }
            sum_red > sum_blue
        }
        else if state.turn() == Color::Red {
            state.nexts().into_iter().any(|s1| self.wins(s1))
        }
        else {
            !state.nexts().into_iter().any(|s1| !self.wins(s1))
        }
    }
    
    fn wins(&mut self, s: u32) -> bool {
        if self.win_vec[s as usize] == 0 {
            let b = self.wins_core(s);
            self.win_vec[s as usize] = if b { 1 } else { 2 }
        }
        self.win_vec[s as usize] == 1
    }
    
    fn read_row() -> [i64; 3] {
        let v: Vec<i64> = read_vec();
        [v[0], v[1], v[2]]
    }
    
    fn read() -> Board {
        let A: [[i64; 3]; 3] = [Board::read_row(),
                                Board::read_row(),
                                Board::read_row()];
        let N = 3usize.pow(9);
        let win_vec: Vec<i32> = vec![0; N];
        Board { A, win_vec }
    }
}


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

fn F(mut board: Board) -> bool {
    board.wins(0)
}

fn main() {
    let board = Board::read();
    let b = F(board);
    println!("{}", if b { "Takahashi" } else { "Aoki" })
}

AtCoder Beginner Contest 349 D

https://atcoder.jp/contests/abc349/tasks/abc349_d

Rを超えないようにできるだけ飛べばよいです。

// Divide Interval
#![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()
}


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

fn read_input() -> (u64, u64) {
    let v = read_vec();
    let (L, R) = (v[0], v[1]);
    (L, R)
}

fn print_intervals(intervals: Vec<(u64, u64)>) {
    println!("{}", intervals.len());
    for (l, r) in intervals.into_iter() {
        println!("{} {}", l, r)
    }
}

fn next(l: u64, R: u64) -> u64 {
    let mut d: u64 = 1;
    loop {
        if l + d > R {
            d >>= 1;
            break
        }
        else if l % d != 0 {
            d >>= 1;
            break
        }
        else {
            d <<= 1
        }
    }
    l + d
}

fn F(L: u64, R: u64) -> Vec<(u64, u64)> {
    let mut intervals: Vec<(u64, u64)> = vec![];
    let mut l = L;
    while l < R {
        let r = next(l, R);
        intervals.push((l, r));
        l = r
    }
    intervals
}

fn main() {
    let (L, R) = read_input();
    let intervals = F(L, R);
    print_intervals(intervals)
}

AtCoder Beginner Contest 348 D

https://atcoder.jp/contests/abc348/tasks/abc348_d

ふつうにQueueを使ってたどればよいです。

// Medicines on Grid
#![allow(non_snake_case)]

use std::collections::VecDeque;


//////////////////// 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()
}

fn YesNo(b: bool) -> String {
    return if b { "Yes" } else { "No" }.to_string()
}


//////////////////// Table ////////////////////

type Point = (usize, usize);

struct Table {
    H: usize,
    W: usize,
    cells: Vec<Vec<(char, i32)>>
}

impl Table {
    fn is_obstacle(&self, pt: Point) -> bool {
        self.cells[pt.0][pt.1].0 == '#'
    }
    
    fn get_medicine(&self, pt: Point) -> i32 {
        self.cells[pt.0][pt.1].1
    }
    
    fn nexts(&self, pt: Point) -> Vec<Point> {
        let (x, y) = pt;
        let mut pts: Vec<Point> = vec![];
        if x > 0 && !self.is_obstacle((x-1, y)) {
            pts.push((x-1, y))
        }
        if x < self.H - 1 && !self.is_obstacle((x+1, y)) {
            pts.push((x+1, y))
        }
        if y > 0 && !self.is_obstacle((x, y-1)) {
            pts.push((x, y-1))
        }
        if y < self.W - 1 && !self.is_obstacle((x, y+1)) {
            pts.push((x, y+1))
        }
        pts
    }
    
    fn find(&self, c: char) -> Option<Point> {
        for (x, v) in self.cells.iter().enumerate() {
            for (y, c1) in v.iter().enumerate() {
                if c1.0 == c {
                    return Some((x, y))
                }
            }
        }
        None
    }
    
    fn read() -> Table {
        let v: Vec<usize> = read_vec();
        let (H, W) = (v[0], v[1]);
        let mut cells: Vec<Vec<(char, i32)>> = (0..H).map(|_| 
                                                        read::<String>().
                                                        chars().
                                                        map(|c| (c, 0)).
                                                        collect::<Vec<_>>()
                                                    ).collect();
        
        let N: usize = read();
        for _ in 0..N {
            let w: Vec<usize> = read_vec();
            let (R, C) = (w[0]-1, w[1]-1);
            let E = w[2] as i32;
            cells[R][C].1 = E
        }
        
        Table { H, W, cells }
    }
}


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

fn F(table: Table) -> bool {
    let mut energies = vec![vec![0; table.W]; table.H];
    let mut queue: VecDeque<(Point, i32)> = VecDeque::new();
    let S = table.find('S').unwrap();
    let T = table.find('T').unwrap();
    queue.push_back((S, table.get_medicine(S)));
    while let Some((pt, energy)) = queue.pop_front() {
        if energy == 0 || energy < energies[pt.0][pt.1] {
            continue
        }
        for pt1 in table.nexts(pt).into_iter() {
            if pt1 == T {
                return true
            }
            let m = table.get_medicine(pt1);
            let e = std::cmp::max(m, energy - 1);
            if e > energies[pt1.0][pt1.1] {
                energies[pt1.0][pt1.1] = e;
                queue.push_back((pt1, e))
            }
        }
    }
    false
}

fn main() {
    let table = Table::read();
    println!("{}", YesNo(F(table)))
}

MojoでProject Euler 27

https://projecteuler.net/problem=27

素数の判定をどれくらい大きさまで行っていいか分からないので、ナイーブに判定してメモ化してみました。
MojoではDictを使えるようになりました。

from collections import Dict

fn main() raises:
    var d = Dict[String, Int]()
    d["a"] = 1
    print(d["a"])       # 1
    print("a" in d)     # True

なので、これをメモ化に使おうとしましたが、

from collections import Dict

var memo = Dict[String, Int]()
fn f(s: String) raises -> Int:
    if s in memo:
        return memo[s]
    elif s == "a":
        memo[s] = 1
        return 1
    else:
        memo[s] = 2
        return 2

fn main() raises:
    print(f("a"))
    print(f("b"))
    print(f("a"))
    print(f("b"))

このテストコードはなぜか落ちます。どうも現状Dictをグローバルな変数にはできないっぽいです。仕方なく、DynamicVectorを使いました。

from collections import Dict
from math import min
import sys


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

def 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)


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

fn calc_min_b(a: Int, L: Int) -> Int:
    if a > 0:
        return 2
    elif -a < L * 2:
        var n = a // 2
        return 2 - (n * n + a * n)
    else:
        return 2 - (L * L + a * L)

var primes = DynamicVector[Int]()
def is_prime(n: Int) -> Bool:
    if n < 2:
        return False
    elif n < primes.size and primes[n] != -1:
        return primes[n] == 1
    
    var b = True
    for p in range(2, n+1):
        if p * p > n:
            break
        elif n % p == 0:
            b = False
            break
    
    if n >= primes.size:
        for m in range(primes.size, n+1):
            primes.push_back(-1)
    primes[n] = 1 if b else 0
    return b

def length(a: Int, b: Int) -> Int:
    var n = 0
    while True:
        m = n * n + a * n + b
        if not is_prime(m):
            break
        n += 1
    return n

fn f(N: Int) raises -> Int:
    var max_length = 0
    var max_product = 0
    for a in range(-N+1, N):
        var min_b = calc_min_b(a, max_length)
        for b in range(min_b, N+1):
            var l = length(a, b)
            if l > max_length:
                max_length = l
                max_product = a * b
                print(a, b, l)
    return max_product

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