AtCoder Beginner Contest 363 D

https://atcoder.jp/contests/abc363/tasks/abc363_d

桁数を指定すれば、その中で何番目の回文数は何かわかります。
オーバーフローの問題はi128を使えばよいです。

// Palindromic Number
#![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()
}

// ex) 90 -> 09
fn reverse_number(mut n: i128) -> i128 {
    let mut m: i128 = 0;
    while n > 0 {
        let d = n % 10;
        m = m * 10 + d;
        n /= 10
    }
    m
}


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

fn num_of_palindromes(e: u32) -> i128 {
    if e == 0 {
        1
    }
    else {
        9 * 10_i128.pow((e-1)/2)
    }
}

fn nth_palindromes(n: i128, e: u32) -> i128 {
    if e == 0 {
        0
    }
    else if e % 2 == 0 {
        let m = n + 10_i128.pow(e/2-1);
        m * 10_i128.pow(e/2) + reverse_number(m)
    }
    else {
        let m = n + 10_i128.pow(e/2);
        m / 10 * 10_i128.pow(e/2+1) + reverse_number(m)
    }
}

fn F(mut N: i128) -> i128 {
    for e in 0.. {
        let np = num_of_palindromes(e);
        if np >= N {
            return nth_palindromes(N - 1, e)
        }
        N -= np
    }
    return 0
}

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

MojoでProject Euler 46

https://projecteuler.net/problem=46

自然数nに対し、 n-2m^2素数なのか調べます。素数のSetを作っておけばいいのですが、答えがNなら、計算量は O(N^{3/2}\log{N})くらいなので、計算量を 10^9くらいに抑えたければ600000くらいまで素数を作っておけばよいです。

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


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

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


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

fn f() -> Int:
    var N = 6*10**5
    var primes = make_prime_table(N)
    var set_primes = Set(primes)
    for n in range(9, N+1, 2):
        if n in primes:
            continue
        for m in range(1, n):
            var s = 2 * m * m
            if s >= n - 1:
                return n
            var p = n - s
            if p in set_primes:
                break
    else:
        return 0

fn main():
    print(f())

MojoでProject Euler 45

https://projecteuler.net/problem=45

 T_k = P_l = H_mとします。まず T_k = H_mは簡単で、
 k(k+1)/2 = m(2m-1)
両辺を8倍して1を足すと、
 (2k+1)^2 = (4m-1)^2
 (2k+4m)(2k-4m+2) = 0
だから、
 k = 2m - 1
です。
 T_k = P_lは、
両辺を24倍して3を足すと、
 3(2k + 1)^2 = (6l-1)^2 + 2
 (6l-1)^2 - 3(2k+1)^2 = -2
 x \equiv 6l-1 \ \ y \equiv 2k+1とおくと、
 x^2 - 3y^2 = -2となって、これはPell方程式を拡張したものです。
この解は、
 x_n + y_n\sqrt{3} = (1+\sqrt{3})(2+\sqrt{3})^n
と書けます。

import sys

fn f() -> Int:
    fn mul(x: Tuple[Int, Int], y: Tuple[Int, Int]) -> Tuple[Int, Int]:
        var a = x.get[0, Int]()
        var b = x.get[1, Int]()
        var c = y.get[0, Int]()
        var d = y.get[1, Int]()
        return (a * c + b * d * 3, a * d + b * c)
    
    var counter = 0
    var x = (1, 1)
    while True:
        x = mul(x, (2, 1))
        var a = x.get[0, Int]()
        var b = x.get[1, Int]()
        if a % 6 == 5 and b % 4 == 3:
            counter += 1
            if counter == 3:
                var l = (a + 1) // 6
                return l * (3 * l - 1) // 2

fn main():
    print(f())

MojoでProject Euler 44

https://projecteuler.net/problem=44

 P_i + P_j = P_kとすると、
 \displaystyle \frac{i(3i-1)}{2} + \frac{j(3j-1)}{2} = \frac{k(3k-1)}{2}
両辺を24倍して2を足すと、
 (6i-1)^2 + (6j-1)^2 = (6k-1)^2 + 1
 x \equiv 6i-1 \ \ y \equiv 6j-1 \ \ z \equiv 6k-1
とすると、
 (x-1)(x+1) = (z-y)(z+y)
だから、ふるいで素因数分解をしておいて、xの小さいほうから(x-1)(x+1)を素因数分解を求めます。そうすると約数が簡単に求められるので、そのときに、yとzが6で割って5余るならとりあえず一方の条件は満たします。そして足しても五角数かを調べます。

import sys

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

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 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 __imul__(inout self, other: Factors):
        for i in range(len(other.ps)):
            var p = other.ps[i]
            var e = other.es[i]
            for j in range(len(self.ps)):
                if self.ps[j] == p:
                    self.es[j] += e
                    break
            else:
                self.ps.append(p)
                self.es.append(e)
            self.value *= p**e
    
    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)
    
    @staticmethod
    fn make_partial_factors_table(first: Int, last: Int,
                                    primes: List[Int]) -> List[Factors]:
        fn mul(inout fs: Factors, p: Int, e: Int):
            fs.ps.append(p)
            fs.es.append(e)
            fs.value *= p**e
        
        var N = last - first
        var a = List[Int](capacity=N)
        for n in range(first, last):
            a.append(n)
        var b = initialize_list(N, Factors.create(1))
        for p in primes:
            if p[] * p[] >= last:
                break
            var n0 = (first + p[] - 1) // p[] * p[]
            for n in range(n0, last, p[]):
                var t = div_pow(a[n-first], p[])
                var e = t.get[0, Int]()
                a[n-first] = t.get[1, Int]()
                mul(b[n-first], p[], e)
        
        for k in range(last - first):
            if a[k] != 1:
                mul(b[k], a[k], 1)
        return b
    
    @staticmethod
    fn divisors_core(ps: List[Int], es: List[Int]) -> List[Int]:
        if len(es) == 1:
            var ds = List[Int](1)
            var p = ps[0]
            var n = 1
            for _ in range(es[0]):
                n *= p
                ds.append(n)
            return ds
        else:
            var mid = len(ps) // 2
            var ps1 = sublist(ps, 0, mid)
            var es1 = sublist(es, 0, mid)
            var ps2 = sublist(ps, mid, len(ps))
            var es2 = sublist(es, mid, len(ps))
            var ds1 = Factors.divisors_core(ps1, es1)
            var ds2 = Factors.divisors_core(ps2, es2)
            var ds = List[Int](capacity=len(ds1)*len(ds2))
            for d1 in ds1:
                for d2 in ds2:
                    ds.append(d1[] * d2[])
            return ds
    
    @staticmethod
    fn divisors(fs: Factors) -> List[Int]:
        if len(fs) == 0:
            return List[Int](1)
        else:
            return Factors.divisors_core(fs.ps, fs.es)


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

fn is_pentagonal(n: Int) -> Bool:
    var m = 1 + 24 * n
    var r = int_sqrt(m)
    if r * r != m:
        return False
    else:
        return (1 + r) % 6 == 0

# n(3n-1)/2 * 24 + 1 = (6n-1)^2
# x^2 + y^2 = z^2 + 1
# x^2 - 1 = z^2 - y^2
fn f() -> Int:
    var primes = make_prime_table(10000)
    var M = 6000
    var first = 4
    while True:
        var fss = Factors.make_partial_factors_table(first, first+M+1, primes)
        for x in range(first+1, first+M-1, 6):
            var n = (x - 1) * (x + 1)
            var fs = fss[x-first-1] * fss[x-first+1]
            var ds = Factors.divisors(fs)
            for d1 in ds:
                var d2 = n // d1[]
                if d2 <= d1[] or (d2 - d1[]) % 2 == 1:
                    continue
                var y = (d2 - d1[]) // 2
                var z = (d1[] + d2) // 2
                if y % 6 == 5 and z % 6 == 5:
                    var P2 = (y * y - 1) // 24
                    var P3 = (z * z - 1) // 24
                    var P4 = P2 + P3
                    if is_pentagonal(P4):
                        return P3 - P2
        first += M

fn main():
    print(f())

MojoでProject Euler 43

https://projecteuler.net/problem=43

素数が大きいほうから絞るとよいと思います。

import sys

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

fn digits(owned n: Int, m: Int) -> List[Int]:
    var ds = List[Int]()
    for _ in range(m):
        var d = n % 10
        ds.append(d)
        n //= 10
    return ds


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

fn unused_digit(used: Int) -> Int:
    for d in range(10):
        if ((used >> d) & 1) == 0:
            return d
    else:
        return -1

fn g(used: Int, n: Int, m: Int, k: Int, primes: List[Int]) -> Int:
    if k == -1:
        var d = unused_digit(used)
        var n1 = d * 10**9 + n
        print(n1)
        return n1
    
    var s = 0
    var p = primes[k]
    for d in range(10):
        if ((used >> d) & 1) == 1:
            continue
        var m1 = d * 100 + m
        if m1 % p == 0:
            var n1 = d * 10**(8-k) + n
            var used1 = used | (1 << d)
            s += g(used1, n1, m1//10, k-1, primes)
    return s

fn f() -> Int:
    var primes = List[Int](2, 3, 5, 7, 11, 13, 17)
    var s = 0
    for k in range(1, 999//17+1):
        var n = 17 * k
        var ds = digits(n, 3)
        var used = 0
        for d in ds:
            if ((used >> d[]) & 1) == 1:
                break
            used |= 1 << d[]
        else:
            s += g(used, n, n // 10, 5, primes)
    return s

fn main():
    print(f())

AtCoder Beginner Contest 362 E

https://atcoder.jp/contests/abc362/tasks/abc362_e

DPです。状態は(次の値, 公差, 長さ)、値は頻度にすればよいですが、次の値ですぐに検索できるようにHashMapのキーを次の値にしておきます。

// Count Arithmetic Subsequences
#![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()
}

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


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

const D: usize = 998244353;

fn read_input() -> Vec<i64> {
    let _N: usize = read();
    let A: Vec<i64> = read_vec();
    A
}

fn F(A: Vec<i64>) {
    let N = A.len();
    // { next: { (d, len): freq } }
    let mut dp: HashMap<i64, HashMap<(i64, usize), usize>> = HashMap::new();
    for i in 1..N {
        let mut v: Vec<(i64, usize, usize)> = vec![];   // [(d, len, freq)]
        if let Some(dic) = dp.get(&A[i]) {
            for (&(d1, len1), &freq1) in dic.iter() {
                v.push((d1, len1+1, freq1))
            }
        }
        for j in 0..i {
            v.push((A[i] - A[j], 2, 1))
        }
        
        for (d, len, freq) in v.into_iter() {
            let n = A[i] + d;
            let e1 = dp.entry(n).or_insert(HashMap::new());
            let e2 = (*e1).entry((d, len)).or_insert(0);
            *e2 = (*e2 + freq) % D;
        }
    }
    
    let mut B: Vec<usize> = vec![0; N+1];
    B[1] = N;
    for (_, dic) in dp.into_iter() {
        for ((_, len), freq) in dic.into_iter() {
            B[len] = (B[len] + freq) % D
        }
    }
    
    println!("{}", (1..N+1).map(|i| B[i].to_string()).
                                    collect::<Vec<String>>().join(" "))
}

fn main() {
    let A = read_input();
    F(A)
}

AtCoder Beginner Contest 362 D

https://atcoder.jp/contests/abc362/tasks/abc362_d

ほぼダイクストラ法ですね。

// Shortest Path 3
#![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 Weight = i64;
type Edge = (Node, Node, Weight);
type Graph = Vec<Vec<(Node, Weight)>>;

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

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


//////////////////// BinaryHeap ////////////////////

use std::collections::BinaryHeap;


#[derive(Debug, Eq, PartialEq)]
struct Item {
    weight: Weight,
    node: Node,
}

impl Ord for Item {
    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
        other.weight.cmp(&self.weight)
    }
}

impl PartialOrd for Item {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}


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

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

fn F(A: Vec<Weight>, edges: Vec<Edge>) {
    let N = A.len();
    let graph = create_graph(N, edges);
    let mut weights: Vec<Weight> = vec![-1; N];
    let mut tree = BinaryHeap::<Item>::new();
    tree.push(Item { weight: A[0], node: 0 });
    while let Some(item) = tree.pop() {
        if weights[item.node] != -1 {
            continue
        }
        
        weights[item.node] = item.weight;
        for &(v, w) in graph[item.node].iter() {
            if weights[v] == -1 {
                tree.push(Item { weight: item.weight + w + A[v], node: v })
            }
        }
    }
    
    println!("{}", (1..N).map(|i| weights[i].to_string()).
                                        collect::<Vec<String>>().join(" "))
}

fn main() {
    let (A, edges) = read_input();
    F(A, edges)
}