同じ場所を何度も通ってそこにたくさんの点があると時間がかかるので、重複を除いてあらかじめ通る場所を計算しておきます。そのために、各固定座標に対し範囲をBTreeSetにし、そこに範囲を追加して、範囲が結合できるなら結合するようにします。

// Santa Claus 2
#![allow(non_snake_case)]

use std::collections::{HashMap, BTreeSet};


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


//////////////////// Point ////////////////////

type Point = (i64, i64);

fn read_point() -> Point {
    let v: Vec<i64> = read_vec();
    (v[0], v[1])
}


//////////////////// Range ////////////////////

type Range = (i64, i64);

fn add_range(rng: Range, ranges: &mut BTreeSet<Range>) {
    let mut to_remove = Vec::new();
    let mut start = rng.0;
    let mut end = rng.1;
    
    for &range in ranges.range((rng.0, i64::MIN)..(rng.1, i64::MAX)) {
        if range.1 < rng.0 || range.0 > rng.1 {
            continue
        }
        to_remove.push(range);
        start = start.min(range.0);
        end = end.max(range.1)
    }
    
    // 影響を受ける範囲を削除
    for range in to_remove {
        ranges.remove(&range);
    }
    
    // 新しいマージされた範囲を追加
    ranges.insert((start, end));

}


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

fn read_input() -> (usize, Point, Vec<Point>) {
    let v: Vec<i64> = read_vec();
    let (N, M) = (v[0] as usize, v[1] as usize);
    let S = (v[2], v[3]);
    let houses: Vec<Point> = (0..N).map(|_| read_point()).collect();
    (M, S, houses)
}

fn classify(keys: &Vec<i64>, values: &Vec<i64>) -> HashMap<i64, BTreeSet<i64>> {
    let mut c: HashMap<i64, BTreeSet<i64>> = HashMap::new();
    for (&k, &v) in keys.iter().zip(values.iter()) {
        let e = c.entry(k).or_insert(BTreeSet::new());
        (*e).insert(v);
    }
    c
}

fn read_move() -> (char, i64) {
    let v: Vec<String> = read_vec();
    let D: char = v[0].chars().next().unwrap();
    let C: i64 = v[1].parse().unwrap();
    (D, C)
}

fn collect_moves(moves: Vec<(char, i64)>, start: Point) ->
                                    (HashMap<i64, BTreeSet<Range>>,
                                     HashMap<i64, BTreeSet<Range>>, Point) {
    let mut range_x_dir: HashMap<i64, BTreeSet<Range>> = HashMap::new();
    let mut range_y_dir: HashMap<i64, BTreeSet<Range>> = HashMap::new();
    let mut x = start.0;
    let mut y = start.1;
    for (D, C) in moves.into_iter() {
        match D {
            'U' => {
                let e = range_y_dir.entry(x).or_insert(BTreeSet::new());
                add_range((y, y+C+1), &mut *e);
                y += C
            }
            'D' => {
                let e = range_y_dir.entry(x).or_insert(BTreeSet::new());
                add_range((y-C, y+1), &mut *e);
                y -= C
            }
            'L' => {
                let e = range_x_dir.entry(y).or_insert(BTreeSet::new());
                add_range((x-C, x+1), &mut *e);
                x -= C
            }
            _   => {
                let e = range_x_dir.entry(y).or_insert(BTreeSet::new());
                add_range((x, x+C+1), &mut *e);
                x += C
            }
        }
    }
    (range_x_dir, range_y_dir, (x, y))
}

fn count(set_xs: &BTreeSet<i64>, x_ranges: &BTreeSet<Range>) -> Vec<i64> {
    let xs: Vec<i64> = set_xs.iter().map(|&x| x).collect();
    let ranges: Vec<Range> = x_ranges.iter().map(|&r| r).collect();
    let mut v: Vec<i64> = vec![];
    let mut k: usize = 0;
    let mut l: usize = 0;
    while k < xs.len() && l < ranges.len() {
        let x = xs[k];
        let (first, last) = ranges[l];
        if first <= x && x < last {
            v.push(x);
            k += 1
        }
        else if x < first {
            k += 1
        }
        else {
            l += 1
        }
    }
    v
}

fn F(M: usize, S: Point, houses: Vec<Point>) {
    let xs: Vec<i64> = houses.iter().map(|&(x, _)| x).collect();
    let ys: Vec<i64> = houses.iter().map(|&(_, y)| y).collect();
    let collection_by_xs = classify(&xs, &ys);
    let collection_by_ys = classify(&ys, &xs);
    let moves: Vec<(char, i64)> = (0..M).map(|_| read_move()).collect();
    let (range_x_dir, range_y_dir, G) = collect_moves(moves, S);
    
    let mut set: BTreeSet<Point> = BTreeSet::new();
    for (y, x_ranges) in range_x_dir.into_iter() {
        match collection_by_ys.get(&y) {
            Some(set_xs) => {
                let xs1 = count(set_xs, &x_ranges);
                for x in xs1.into_iter() {
                    set.insert((x, y));
                }
            },
            None       => ()
        }
    }
    for (x, y_ranges) in range_y_dir.into_iter() {
        match collection_by_xs.get(&x) {
            Some(set_ys) => {
                let ys1 = count(set_ys, &y_ranges);
                for y in ys1.into_iter() {
                    set.insert((x, y));
                }
            },
            None       => ()
        }
    }
    println!("{} {} {}", G.0, G.1, set.len())
}

fn main() {
    let (M, S, houses) = read_input();
    F(M, S, houses)
}

2024-12-20

MojoでProject Euler 83

https://projecteuler.net/problem=83

ダイクストラ法ですが、実装するのはけっこう大変ですね。TupleがComparableCollectionElementとKeyElementを実装してくれれば簡単なんですけどね。

import sys


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

trait Printable(CollectionElement, Stringable):
    pass

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


#################### HeapQueue ####################

struct HeapQueue[T: ComparableCollectionElement]:
    var a: List[T]
    
    fn __init__(inout self):
        self.a = List[T]()
    
    fn is_empty(self) -> Bool:
        return len(self.a) == 0
    
    fn pop(inout self) -> T:
        var top = self.a[0]
        self.a[0] = self.a.pop()
        var N = self.a.size
        var i = 0
        while i < N-1:
            var j1 = i*2+1
            var j2 = i*2+2
            var j: Int = 0
            if j1 >= N:
                break
            elif j2 == N:
                j = j1
            else:
                if self.a[j1] <= self.a[j2]:
                    j = j1
                else:
                    j = j2
            if self.a[j] < self.a[i]:
                self.swap(i, j)
                i = j
            else:
                break
        
        return top
    
    fn push(inout self, e: T):
        self.a.append(e)
        var i = self.a.size - 1
        while i > 0:
            var j = (i-1) // 2
            if self.a[j] <= self.a[i]:
                break
            else:
                self.swap(i, j)
                i = j
    
    fn swap(inout self, i: Int, j: Int):
        var tmp = self.a[i]
        self.a[i] = self.a[j]
        self.a[j] = tmp^


#################### Point ####################

@value
struct Point(KeyElement):
    var x: Int
    var y: Int
    
    fn __eq__(self, other: Point) -> Bool:
        return self.x == other.x and self.y == other.y
    
    fn __ne__(self, other: Point) -> Bool:
        return self.x != other.x or self.y != other.y
    
    fn __hash__(self) -> Int:
        return self.x * 10000 + self.y


#################### Item ####################

@value
struct Item(Comparable):
    var value: Int
    var pt: Node
    
    fn __lt__(self, other: Item) -> Bool:
        return self.value < other.value
    
    fn __le__(self, other: Item) -> Bool:
        return self.value <= other.value
    
    fn __eq__(self, other: Item) -> Bool:
        return self.value == other.value
    
    fn __ne__(self, other: Item) -> Bool:
        return self.value != other.value
    
    fn __gt__(self, other: Item) -> Bool:
        return self.value > other.value
    
    fn __ge__(self, other: Item) -> Bool:
        return self.value >= other.value


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


#################### Matrix ####################

fn read_matrix(path: String) -> List[List[Int]]:
    try:
        var matrix = List[List[Int]]()
        with open(path, "r") as f:
            var whole = f.read()
            var lines = whole.split('\n')
            for line in lines:
                if line[] == '':
                    break
                var ns = List[Int]()
                var v = line[].split(',')
                for s in v:
                    var n = atol(s[])
                    ns.append(n)
                matrix.append(ns)
        return matrix
    except:
        return List[List[Int]]()


#################### Graph ####################

alias Node = Point

struct Graph:
    var g: Dict[Node, List[Node]]
    var table: List[List[Int]]
    
    fn __init__(inout self, g: Dict[Node, List[Node]], table: List[List[Int]]):
        self.g = g
        self.table = table
    
    fn value(self, pt: Point) -> Int:
        return self.table[pt.x][pt.y]
    
    @staticmethod
    fn create(table: List[List[Int]]) raises -> Graph:
        var H = len(table)
        var W = len(table[0])
        var g = Dict[Node, List[Node]]()
        for x in range(H):
            for y in range(W):
                var pt = Point(x, y)
                g[pt] = List[Node]()
                if x != 0:
                    g[pt].append(Point(x-1, y))
                if x != H - 1:
                    g[pt].append(Point(x+1, y))
                if y != 0:
                    g[pt].append(Point(x, y-1))
                if y != W - 1:
                    g[pt].append(Point(x, y+1))
        return Graph(g, table)


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

fn Dijkstra(graph: Graph, start: Node, goal: Node) raises -> Int:
    var queue = HeapQueue[Item]()
    queue.push(Item(graph.value(start), start))
    var visited = Set[Node]()
    while not queue.is_empty():
        var e = queue.pop()
        visited.add(e.pt)
        if e.pt == goal:
            return e.value
        for dest in graph.g[e.pt]:
            if dest[] not in visited:
                queue.push(Item(e.value + graph.value(dest[]), dest[]))
    return -1

fn f(path: String) -> Int:
    var table = read_matrix(path)
    var H = len(table)
    var W = len(table[0])
    try:
        var graph = Graph.create(table)
        return Dijkstra(graph, Point(0, 0), Point(H-1, W-1))
    except:
        return 0

fn main():
    var path = String("0083_matrix.txt")
    print(f(path))

2024-12-19

MojoでProject Euler 82

https://projecteuler.net/problem=82

上下に動けるので一気にDPを使えず、マスごとに上から来るのと下から来るのと左から直接来るのとで最も小さいものを選びます。

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


#################### Matrix ####################

fn read_matrix(path: String) -> List[List[Int]]:
    try:
        var matrix = List[List[Int]]()
        with open(path, "r") as f:
            var whole = f.read()
            var lines = whole.split('\n')
            for line in lines:
                if line[] == '':
                    break
                var ns = List[Int]()
                var v = line[].split(',')
                for s in v:
                    var n = atol(s[])
                    ns.append(n)
                matrix.append(ns)
        return matrix
    except:
        return List[List[Int]]()

trait Printable(CollectionElement, Stringable):
    pass

fn print_matrix[T: Printable](mat: List[List[T]]):
    var s = String('[\n')
    for v in mat:
        s += ' [' + str(v[][0])
        for i in range(1, len(v[])):
            s += ', ' + str(v[][i])
        s += ']\n'
    s += '\n]'
    print(s)

fn initialize_matrix[T: CollectionElement](a: T,
                                    H: Int, W: Int) -> List[List[T]]:
    var matrix = List[List[T]]()
    for _ in range(H):
        var v = List[T]()
        for _ in range(W):
            v.append(a)
        matrix.append(v)
    return matrix


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

fn initialize_dp(mat: List[List[Int]]) -> List[Int]:
    var dp = List[Int]()
    for v in mat:
        dp.append(v[][0])
    return dp

fn min_path(i: Int, j: Int, mat: List[List[Int]], dp: List[Int]) -> Int:
    var H = len(dp)
    var a = initialize_list(H, 0)
    for i1 in range(i):
        if i1 == 0:
            a[0] = dp[0] + mat[0][j]
        else:
            a[i1] = min(dp[i1], a[i1-1]) + mat[i1][j]
    for i1 in range(H-1, i, -1):
        if i1 == H-1:
            a[i1] = dp[i1] + mat[i1][j]
        else:
            a[i1] = min(dp[i1], a[i1+1]) + mat[i1][j]
    
    if i == 0:
        return min(dp[i], a[i-1]) + mat[i][j]
    elif i == H-1:
        return min(dp[i], a[i+1]) + mat[i][j]
    else:
        return min(dp[i], min(a[i-1], a[i+1])) + mat[i][j]

fn update_dp(j: Int, mat: List[List[Int]], dp: List[Int]) -> List[Int]:
    var new_dp = List[Int]()
    var H = len(dp)
    for i in range(H):
        new_dp.append(min_path(i, j, mat, dp))
    return new_dp

fn f(path: String) -> Int:
    var mat = read_matrix(path)
    var H = len(mat)
    var W = len(mat[0])
    var dp = initialize_dp(mat)
    for j in range(1, W):
        dp = update_dp(j, mat, dp)
    
    var min_value = dp[0]
    for i in range(1, H):
        min_value = min(min_value, dp[i])
    return min_value

fn main() raises:
    var path = String("0082_matrix.txt")
    print(f(path))

2024-12-18

MojoでProject Euler 81

https://projecteuler.net/problem=81

ループがない有向グラフなので、DPを使えばよいです。

import sys


#################### Matrix ####################

fn read_matrix(path: String) -> List[List[Int]]:
    try:
        var matrix = List[List[Int]]()
        with open(path, "r") as f:
            var whole = f.read()
            var lines = whole.split('\n')
            for line in lines:
                if line[] == '':
                    break
                var ns = List[Int]()
                var v = line[].split(',')
                for s in v:
                    var n = atol(s[])
                    ns.append(n)
                matrix.append(ns)
        return matrix
    except:
        return List[List[Int]]()

trait Printable(CollectionElement, Stringable):
    pass

fn print_matrix[T: Printable](mat: List[List[T]]):
    var s = String('[\n')
    for v in mat:
        s += ' [' + str(v[][0])
        for i in range(1, len(v[])):
            s += ', ' + str(v[][i])
        s += ']\n'
    s += '\n]'
    print(s)

fn initialize_matrix[T: CollectionElement](a: T,
                                    H: Int, W: Int) -> List[List[T]]:
    var matrix = List[List[T]]()
    for _ in range(H):
        var v = List[T]()
        for _ in range(W):
            v.append(a)
        matrix.append(v)
    return matrix


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

fn f(path: String) -> Int:
    var mat = read_matrix(path)
    var H = len(mat)
    var W = len(mat[0])
    var dp = initialize_matrix(10**18, H+1, W+1)
    dp[0][1] = 0
    dp[1][0] = 0
    for i in range(1, H+1):
        for j in range(1, W+1):
            dp[i][j] = min(dp[i-1][j], dp[i][j-1]) + mat[i-1][j-1]
    return dp[H][W]

fn main() raises:
    var path = String("0081_matrix.txt")
    print(f(path))

2024-12-17

AtCoder Beginner Contest 384 F

https://atcoder.jp/contests/abc384/tasks/abc384_f

ふつうに計算すれば $O(N^2)$ 以上かかりますが、Aを偶数と奇数に分けてみましょう。そうすると、偶数から一つ奇数から一つ取って足せば割らなくてもいいので、和を簡単に計算できます。これがd=2のときです。d=4のときは、4で割った余りで分類して、余りを足して2または6になるペアで簡単に和が取れます。これをd=8、16と処理していき、dになるペアがなくなったら終わりです。

// Double Sum 2
#![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 ////////////////////

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

fn div_2pow(mut n: i64) -> i64 {
    while n % 2 == 0 {
        n /= 2
    }
    n
}

fn sum_same(v: &Vec<i64>) -> i64 {
    v.iter().map(|&e| e).sum::<i64>() * (v.len() as i64 - 1)
}

fn sum_diff(v: &Vec<i64>, w: &Vec<i64>) -> i64 {
    v.iter().map(|&e| e).sum::<i64>() * (w.len() as i64) +
    w.iter().map(|&e| e).sum::<i64>() * (v.len() as i64)
}

fn classify(v: &Vec<i64>, d: i64) -> HashMap<i64, Vec<i64>> {
    let mut dic: HashMap<i64, Vec<i64>> = HashMap::new();
    for &n in v.iter() {
        let r = n & (d - 1);
        let e = dic.entry(r).or_insert(vec![]);
        (*e).push(n)
    }
    dic
}

fn is_remained(dic: &HashMap<i64, Vec<i64>>, d: i64) -> bool {
    for (&r1, ns1) in dic.into_iter() {
        let r2 = if r1 == 0 { 0 } else { d - r1 };
        if r1 == r2 {
            if ns1.len() > 1 {
                return true
            }
        }
        else if r1 < r2 {
            if dic.contains_key(&r2) {
                return true
            }
        }
    }
    false
}

fn F(A: Vec<i64>) -> i64 {
    let mut s = A.iter().map(|&n| div_2pow(n)).sum::<i64>();
    let mut d: i64 = 2;
    loop {
        let h = d / 2;
        let dic = classify(&A, d);
        for (&r1, ns1) in dic.iter() {
            let r2 = if r1 <= h { h - r1 } else { h * 3 - r1 };
            if r1 == r2 {
                s += sum_same(ns1) / h
            }
            else if r1 < r2 {
                match dic.get(&r2) {
                    Some(ns2) => { s += sum_diff(ns1, ns2) / h },
                    None      => ()
                }
            }
        }
        
        if !is_remained(&dic, d) {
            break
        }
        
        d <<= 1
    }
    s
}

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

2024-12-16

AtCoder Beginner Contest 384 E

https://atcoder.jp/contests/abc384/tasks/abc384_e

周りのマスをQueueに入れて、値が小さい順に取り出します。

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


//////////////////// Matirx ////////////////////

type Point = (usize, usize);
type Matrix = Vec<Vec<u64>>;

fn read_matrix(H: usize) -> Matrix {
    (0..H).map(|_| read_vec::<u64>()).collect::<Vec<_>>()
}


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

use std::collections::BinaryHeap;

#[derive(Debug, Eq, PartialEq)]
struct Item {
    value: u64,
    pt: Point
}

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

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


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

fn read_input() -> (u64, usize, usize, Matrix) {
    let v: Vec<usize> = read_vec();
    let (H, X) = (v[0], v[2] as u64);
    let w: Vec<usize> = read_vec();
    let (P, Q) = (w[0] - 1, w[1] - 1);
    let S = read_matrix(H);
    (X, P, Q, S)
}

fn neighbors(pt: Point, S: &Matrix) -> Vec<Point> {
    let H = S.len();
    let W = S[0].len();
    let (x, y) = pt;
    let mut neighs: Vec<Point> = vec![];
    if x > 0 {
        neighs.push((x - 1, y))
    }
    if x < H - 1 {
        neighs.push((x + 1, y))
    }
    if y > 0 {
        neighs.push((x, y - 1))
    }
    if y < W - 1 {
        neighs.push((x, y + 1))
    }
    neighs
}

fn F(X: u64, P: usize, Q: usize, S: Matrix) -> u64 {
    let H = S.len();
    let W = S[0].len();
    let mut queued: Vec<Vec<bool>> = vec![vec![false; W]; H];
    queued[P][Q] = true;
    let mut s = S[P][Q];
    let mut heap = BinaryHeap::new();
    for pt in neighbors((P, Q), &S).into_iter() {
        heap.push(Item { value: S[pt.0][pt.1], pt });
        queued[pt.0][pt.1] = true
    }
    
    while let Some(item) = heap.pop() {
        if item.value >= (s + X - 1) / X {
            break
        }
        s += item.value;
        for pt in neighbors(item.pt, &S) {
            if !queued[pt.0][pt.1] {
                heap.push(Item { value: S[pt.0][pt.1], pt });
                queued[pt.0][pt.1] = true
            }
        }
    }
    s
}

fn main() {
    let (X, P, Q, S) = read_input();
    println!("{}", F(X, P, Q, S))
}

2024-12-15

AtCoder Beginner Contest 384 D

https://atcoder.jp/contests/abc384/tasks/abc384_d

繰り返し部分を除いて、数列Aの中か両端で和がSになる範囲を探せばいいですが、尺取り法でO(N)でできます。

// Repeated Sequence
#![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()
}

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


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

fn read_input() -> (i64, Vec<i64>) {
    let v: Vec<i64> = read_vec();
    let S = v[1];
    let A: Vec<i64> = read_vec();
    (S, A)
}

fn accumulate(A: &Vec<i64>) -> Vec<i64> {
    let mut acc: Vec<i64> = vec![0];
    for &a in A.iter() {
        acc.push(acc[acc.len()-1] + a)
    }
    acc
}

fn is_equal_mid(S: i64, acc: &Vec<i64>) -> bool {
    let mut k = 0;
    let mut l = 0;
    while l < acc.len() {
        if acc[l] - acc[k] == S {
            return true
        }
        else if acc[l] - acc[k] < S {
            l += 1
        }
        else {
            k += 1
        }
    }
    false
}

fn is_equal_edge(S: i64, acc: &Vec<i64>) -> bool {
    is_equal_mid(*acc.iter().last().unwrap() - S, acc)
}

fn F(S: i64, A: Vec<i64>) -> bool {
    let acc = accumulate(&A);
    let sum = *acc.iter().last().unwrap();
    let r = S % sum;
    is_equal_mid(r, &acc) || is_equal_edge(r, &acc)
}

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

inamori’s diary

AtCoder Beginner Contest 385 D