LoginSignup
37
48

More than 1 year has passed since last update.

Rustで競技プログラミングよくばりセット

Last updated at Posted at 2022-09-13

競技プログラミングを行う上で使用頻度の高いアルゴリズムやデータ構造を、Rust(プログラミング言語)で実装するためのノートです。各アルゴリズムの実装と、それを用いる問題や解答例を掲載しています。

目次

関数/メソッド

データ構造

アルゴリズム

グラフ関連

整数論関連

関数/メソッド

最小値/最大値 (min/max)

  • 最小値/最大値を取るための関数です。
  • 標準のstd::cmp::{min, max}で実現できます。
use std::cmp;

fn main() {
    // min
    println!("{}", cmp::min(1, 2)); // 1
    println!("{}", cmp::min(1, -1)); // -1

    // max
    println!("{}", cmp::max(1, 2)); // 2
    println!("{}", cmp::max(1, -1)); // 1
}

  1. AtCoder Beginner Contest 052 A - Two Rectangles【問題 / 解答例
  2. AtCoder Beginner Contest 098 A - Add Sub Mul【問題 / 解答例

絶対値 (abs)

  • 絶対値を取るためのメソッドです。
  • 標準のabs()で使えます。
fn main (){
    println!("{}", (-1isize).abs()); // 1
    println!("{}", (1isize).abs()); // 1
    println!("{}", (-100isize).abs()); // 100

    // !注意! 次の例では-100になる。 (∵ -100isize.abs() -> -(100isize.abs()) -> -100 のように解釈されるため)
    println!("{}", -100isize.abs()); // -100
}

  1. AtCoder Beginner Contest 188 A - Three-Point Shot【問題 / 解答例

スワップ (swap)

  • 2つの値をスワップする関数です。
  • 標準のstd::mem::swap()でできます。
use std::mem;

fn main (){
    let mut a = 1;
    let mut b = 2;
    println!("a: {}, b: {}", a, b); // a: 1, b: 2
    mem::swap(&mut a, &mut b);
    println!("a: {}, b: {}", a, b); // a: 2, b: 1
}

  1. AtCoder Beginner Contest 161 A - ABC Swap【問題 / 解答例

ソート (sort)

sort()

  • Vectorの要素を特定の順番で並び替えるためのメソッドです。
  • 標準のsort()で昇順に並び替えることができます。
fn main() {
    let mut v = vec![3, 1, 4, 1, 5, 9, 2];
    v.sort();
    println!("{:?}", v); // [1, 1, 2, 3, 4, 5, 9]
}

  1. Chokudai SpeedRun 001 D - ソート【問題 / 解答例

sort_by()

  • sort_by()では、比較関数をもとにして並び替えることができます。
  • 例えば、cmp()が出力する順序関係を逆転させることで、降順に並び替えることができます。
fn main() {
    let mut v = vec![3, 1, 4, 1, 5, 9, 2];
    v.sort_by(|a, b| a.cmp(b).reverse());
    println!("{:?}", v); // [9, 5, 4, 3, 2, 1, 1]
}

  1. AtCoder Beginner Contest 260 B - Better Students Are Needed!【問題 / 解答例

また、sort_by()partial_cmp()組み合わせることによって、Ordトレイトが実装されていないf64型のVecをソートすることもできます。(Rustの小数にはOrdトレイトが実装されていないため、sort()はできません。)

fn main() {
    let mut v = vec![3.0, 1.0, 4.0, 1.0, 5.0];
    v.sort_by(|a, b| a.partial_cmp(b).unwrap());
    println!("{:?}", v); // [1.0, 1.0, 3.0, 4.0, 5.0]
}

  1. AtCoder Beginner Contest 138 C - Alchemist【問題 / 解答例

sort_by_key()

  • sort_by_key()では、比較関数をもとにして並び替えることができます。
  • 例えば、次のようにVectorをtupleのi番目の要素をもとにして並び替えることができます。
fn main() {
    let mut v = vec![(3, 1, 4), (1, 5, 9), (2, 6, 5)];
    // 各tupleの0番目の要素をもとにソート
    v.sort_by_key(|x| x.0);
    println!("{:?}", v); // [(_1_, 5, 9), (_2_, 6, 5), (_3_, 1, 4)]

    // 各tupleの1番目の要素をもとにソート
    v.sort_by_key(|x| x.1);
    println!("{:?}", v); // [(3, _1_, 4), (1, _5_, 9), (2, _6_, 5)]

    // 各tupleの2番目の要素をもとにソート
    v.sort_by_key(|x| x.2);
    println!("{:?}", v); // [(3, 1, _4_), (2, 6, _5_), (1, 5, _9_)]
}

反転 (reverse)

  • Vectorの要素を反転させるためのメソッドです。
  • 標準のrevese()で反転させることができます。
fn main() {
    let mut v = vec![3, 1, 4, 1, 5, 9, 2];
    v.reverse();
    println!("{:?}", v); // [2, 9, 5, 1, 4, 1, 3]
}

重複要素の削除 (unique, dedup)

  • 連続した要素を削除するためのメソッドです。
  • 標準のdedup()で行うことができます。
fn main() {
    let mut v = vec![1, 1, 2, 3, 4, 5, 5];
    v.dedup();
    println!("{:?}", v); // [1, 2, 3, 4, 5]

    let mut u = vec![1, 1, 1, 2, 2, 1, 1, 1, 0, 3, 3];
    u.dedup();
    println!("{:?}", u); // [1, 2, 1, 0, 3]
}

データ構造

動的配列 (vector)

標準のstd::vecで実現できます。

fn main() {
    let mut v = vec![0, 1, 2, 3, 4];
    println!("{}", v.len()); // 5
    println!("{:?}", v); // [0, 1, 2, 3, 4]

    v[2] =  5;
    println!("{:?}", v); // [0, 1, 5, 3, 4]
}

スタック (stack)

上で述べたstd::vecのpush()とpop()を使うと実現できます。

fn main() {
    let mut stack = vec![];

    stack.push(0);
    stack.push(1);
    stack.push(2);
    println!("{:?}", stack); // [0, 1, 2]

    println!("{:?}", stack.pop()); // Some(2)

    println!("{:?}", stack); // [0, 1]
}

キュー (queue)

標準のstd::collections::VecDequeで実現できます。

use std::collections::VecDeque;

fn main() {
    let mut que = VecDeque::new();

    que.push_back(0);
    que.push_back(1);
    que.push_back(2);
    println!("{:?}", que); // [0, 1, 2]

    println!("{:?}", que.pop_front()); // Some(0)
    println!("{:?}", que); // [1, 2]

    println!("{:?}", que.pop_front()); // Some(1)
    println!("{:?}", que); // [2]

    println!("{:?}", que.pop_front()); // Some(2)
    println!("{:?}", que); // []

    println!("{:?}", que.pop_front()); // None
}

優先度付きキュー (priority queue)

use std::collections::BinaryHeap;

fn main() {
    let mut que = BinaryHeap::new();

    que.push(0);
    que.push(2);
    que.push(1);

    println!("{:?}", que.pop()); // Some(2)
    println!("{:?}", que.pop()); // Some(1)
    que.push(3);
    println!("{:?}", que.pop()); // Some(3)
    println!("{:?}", que.pop()); // Some(0)
    println!("{:?}", que.pop()); // None
}

  1. AtCoder Beginner Contest 137 D - Summer Vacation【問題 / 解答例

Set

標準のstd::collections::HashSetで実現できます。

use std::collections::HashSet;

fn main() {
    let mut set = HashSet::new();

    set.insert(0);
    set.insert(1);
    set.insert(2);

    println!("{}", set.contains(&0)); // true
    println!("{}", set.contains(&3)); // false

    set.remove(&0);
    println!("{}", set.contains(&0)); // false
}

  1. ユニークビジョンプログラミングコンテスト2022(AtCoder Beginner Contest 248) A - Lacked Number【問題 / 解答例

Map

標準のstd::collections::HashMapで実現できます。

use std::collections::HashMap;

fn main() {
    let mut map = HashMap::new();

    map.insert("zero", 0);
    map.insert("one", 1);
    map.insert("two", 2);

    println!("{:?}", map.get("zero")); // Some(0)
    println!("{:?}", map.get("three")); // None
}

Union-Find木 (素集合データ構造)

  • rankの大きさをもとにして併合(union)を行います。
  • 標準のenumであるOrderingで順序を扱えます。
use std::cmp::Ordering;
 
#[derive(Clone)]
pub struct UnionFind {
    parent: Vec<usize>,
    rank: Vec<usize>,
}

impl UnionFind {
    pub fn new(n: usize) -> Self {
        UnionFind {
            parent: (0..n).collect::<Vec<_>>(),
            rank: vec![0; n],
        }
    }

    pub fn is_root(&self, x: usize) -> bool {
        self.parent[x] == x
    }

    pub fn find(&mut self, x: usize) -> usize {
        if self.is_root(x) {
            x
        } else {
            let root = self.find(self.parent[x]);
            self.parent[x] = root;
            root
        }
    }

    pub fn union(&mut self, x: usize, y: usize) {
        let (root_x, root_y) = (self.find(x), self.find(y));

        if root_x != root_y {
            match self.rank[root_x].cmp(&self.rank[root_y]) {
                Ordering::Less => {
                    self.parent[root_x] = root_y;
                }
                Ordering::Equal => {
                    self.parent[root_y] = root_x;
                    self.rank[root_x] += 1;
                }
                Ordering::Greater => {
                    self.parent[root_y] = root_x;
                }
            }
        }
    }

    pub fn is_same(&mut self, x: usize, y: usize) -> bool {
        self.find(x) == self.find(y)
    }
}

  1. AtCoder Typical Contest 001 B - Union Find【問題 / 解答例
  • ある要素が属しているグループに含まれるノードの数を調べることができるように拡張すると、次のようになります。(get_sizeメソッドを追加しています)
use std::cmp::Ordering;

#[derive(Clone)]
pub struct UnionFind {
    parent: Vec<usize>,
    rank: Vec<usize>,
    size: Vec<usize>,
}

impl UnionFind {
    pub fn new(n: usize) -> Self {
        UnionFind { parent: (0..n).collect::<Vec<_>>(), rank: vec![0; n], size: vec![1; n] }
    }

    pub fn is_root(&self, x: usize) -> bool {
        self.parent[x] == x
    }

    pub fn find(&mut self, x: usize) -> usize {
        if self.is_root(x) {
            x
        } else {
            let root = self.find(self.parent[x]);
            self.parent[x] = root;
            root
        }
    }

    pub fn union(&mut self, x: usize, y: usize) {
        let (root_x, root_y) = (self.find(x), self.find(y));

        if root_x != root_y {
            match self.rank[root_x].cmp(&self.rank[root_y]) {
                Ordering::Less => {
                    self.parent[root_x] = root_y;
                    self.size[root_y] += self.size[root_x];
                }
                Ordering::Equal => {
                    self.parent[root_y] = root_x;
                    self.size[root_x] += self.size[root_y];
                    self.rank[root_x] += 1;
                }
                Ordering::Greater => {
                    self.parent[root_y] = root_x;
                    self.size[root_x] += self.size[root_y];
                }
            }
        }
    }

    pub fn is_same(&mut self, x: usize, y: usize) -> bool {
        self.find(x) == self.find(y)
    }

    pub fn get_size(&mut self, x: usize) -> usize {
        if self.is_root(x) {
            self.size[x]
        } else {
            let root = self.find(x);
            self.size[root]
        }
    }
}

  1. AtCoder Beginner Contest 120 D - Decayed Bridges【問題 / 解答例
  2. AtCoder Regular Contest 107 C - Shuffle Permutation【問題 / 解答例

Mod Int (Mint, ModInt)

  • 演算子オーバーロードを用いると、四則演算が自然に行えるようになります。
  • 四則演算の他には、inverse(逆元), pow(冪乗の計算, $a^{n}$), conbinarion(組み合わせ, $nCr$), permutation(順列, $nPr$), factorial(階乗, $n!$)などを実装しています。
use std::ops;

const MOD: usize = 1000000007;

#[derive(Copy, Clone)]
pub struct ModInt {
    value: usize,
}

impl ModInt {
    pub fn new(value: usize) -> ModInt {
        ModInt { value: value % MOD }
    }

    pub fn value(&self) -> usize {
        self.value
    }

    pub fn inverse(&self) -> ModInt {
        // (a, b) -> (x, y) s.t. a * x + b * y = gcd(a, b)
        fn extended_gcd(a: isize, b: isize) -> (isize, isize) {
            if (a, b) == (1, 0) {
                (1, 0)
            } else {
                let (x, y) = extended_gcd(b, a % b);
                (y, x - (a / b) * y)
            }
        }

        let (x, _y) = extended_gcd(self.value() as isize, MOD as isize);
        ModInt::new((MOD as isize + x) as usize)
    }

    // a^n
    pub fn pow(&self, mut n: usize) -> ModInt {
        let mut res = ModInt::new(1);
        let mut x = *self;
        while n > 0 {
            if n % 2 == 1 {
                res *= x;
            }
            x *= x;
            n /= 2;
        }
    
        res
    }
}

impl ops::Add for ModInt {
    type Output = ModInt;
    fn add(self, other: Self) -> Self {
        ModInt::new(self.value + other.value)
    }
}

impl ops::Sub for ModInt {
    type Output = ModInt;
    fn sub(self, other: Self) -> Self {
        ModInt::new(MOD + self.value - other.value)
    }
}

impl ops::Mul for ModInt {
    type Output = ModInt;
    fn mul(self, other: Self) -> Self {
        ModInt::new(self.value * other.value)
    }
}

impl ops::Div for ModInt {
    type Output = ModInt;
    fn div(self, other: Self) -> Self {
        self * other.inverse()
    }
}

impl ops::AddAssign for ModInt {
    fn add_assign(&mut self, other: Self) {
        *self = *self + other;
    }
}

impl ops::SubAssign for ModInt {
    fn sub_assign(&mut self, other: Self) {
        *self = *self - other;
    }
}

impl ops::MulAssign for ModInt {
    fn mul_assign(&mut self, other: Self) {
        *self = *self * other;
    }
}

impl ops::DivAssign for ModInt {
    fn div_assign(&mut self, other: Self) {
        *self = *self / other;
    }
}

// n!
pub fn factorial(n: usize) -> ModInt {
    (1..=n).fold(ModInt::new(1), |x, y| x * ModInt::new(y))
}

// nPr (https://ja.wikipedia.org/wiki/%E9%A0%86%E5%88%97#%E9%A0%86%E5%88%97%E3%81%AE%E6%95%B0%E3%81%88%E4%B8%8A%E3%81%92)
pub fn permutation(n: usize, r: usize) -> ModInt {
    if n < r {
        ModInt::new(0)
    } else {
        (n - r + 1..=n).fold(ModInt::new(1), |x, y| x * ModInt::new(y))
    }
}

// nCr (https://ja.wikipedia.org/wiki/%E7%B5%84%E5%90%88%E3%81%9B_(%E6%95%B0%E5%AD%A6))
pub fn combination(n: usize, r: usize) -> ModInt {
    if n < r {
        ModInt::new(0)
    } else {
        permutation(n, r) / factorial(r)
    }
}

// nHr (https://ja.wikipedia.org/wiki/%E9%87%8D%E8%A4%87%E7%B5%84%E5%90%88%E3%81%9B#%E9%87%8D%E8%A4%87%E7%B5%84%E5%90%88%E3%81%9B%E3%81%AE%E7%B7%8F%E6%95%B0)
pub fn homogeneous(n: usize, r: usize) -> ModInt {
    combination(n + r - 1, r)
}
fn main() {
    let mut a = ModInt::new(1000000000);
    a += ModInt::new(9);
    println!("{}", a.value()); // 2
  
    println!("{}", ModInt::new(2).inverse().value()); // 500000004
    println!("{}", ModInt::new(3).inverse().value()); // 333333336
}

  1. AtCoder Beginner Contest 034 C - 経路【問題 / 解答例
  2. AtCoder Regular Contest 145 C - Split and Maximize 【問題 / 解答例

Segment Tree (セグメント木)

  • RustではACL(AtCoder Library)が使えないので各自で実装をします。

  • 非再帰型の実装です。
    Q5aMu1Td.jpg

  • モノイド$(S, \bullet(op), id)(S:$ 集合、$\bullet: S$上の二項演算、$id: S$の$\bullet$に関する単位元$)$を用いて、数列$\lbrace a_i \rbrace_{i=1}^n$に対して、次の1~3ができます。

    1. update: $a_{i}$ をxに更新
    2. get: $a_{i}$ の取得
    3. fold: ($a_{l} \bullet a_{l+1} \bullet \cdots \bullet a_{r - 1}$) の取得(半開区間$[l, r)$)
  • next_power_of_twoを用いると、セグメント木のノードの数を簡単に決定できます。

  • 具体的な使い方は例の問題を見てください。

pub trait Monoid {
    type S;
    fn op(a: &Self::S, b: &Self::S) -> Self::S;
    fn id() -> Self::S;
}

pub struct SegmentTree<M>
where
    M: Monoid,
{
    size: usize,
    data: Vec<M::S>,
}

impl<M> SegmentTree<M>
where
    M: Monoid,
    M::S: Clone,
{
    /// Creates a segment tree with $\\{a_{i}\\}_{i=1}^{N}$ inside.
    /// n: lenght of $\\{a_{i}\\}_{i=1}^{N}$ (i.e. N)
    pub fn new(n: usize) -> Self {
        let size = n.next_power_of_two();
        SegmentTree::<M> { size, data: vec![M::id(); 2 * size - 1] }
    }

    /// Returns lenght of $\\{a_{i}\\}_{i=1}^{N}$ (i.e. Return N)
    pub fn len(&self) -> usize {
        self.size
    }

    /// Sets x to $a_{idx}$.
    pub fn set(&mut self, mut idx: usize, x: M::S) {
        idx += self.size - 1;
        self.data[idx] = x;
    }

    /// Builds segment tree.
    pub fn build(&mut self) {
        for idx in (0..self.size - 1).rev() {
            self.data[idx] = M::op(&self.data[2 * idx + 1], &self.data[2 * idx + 2]);
        }
    }

    /// Updates $a_{idx}$ to x.
    pub fn update(&mut self, mut idx: usize, x: M::S) {
        idx += self.size - 1;
        self.data[idx] = x;
        while idx > 0 {
            idx = (idx - 1) / 2;
            self.data[idx] = M::op(&self.data[2 * idx + 1], &self.data[2 * idx + 2]);
        }
    }

    /// Returns $a_{idx}$.
    pub fn get(&self, mut idx: usize) -> M::S {
        idx += self.size - 1;
        self.data[idx].clone()
    }

    /// Returns the result (fold op $\left[a_{l}, ... ,a_{r}\right)).$
    /// (i.e. Return $a_{l} (op) a_{l + 1} (op) \cdots (op) a_{r-1})$
    /// Notice that this is a half-opened section.
    pub fn fold(&self, mut l: usize, mut r: usize) -> M::S {
        l += self.size - 1;
        r += self.size - 1;

        let mut xl = M::id();
        let mut xr = M::id();

        while l < r {
            if l % 2 == 0 {
                xl = M::op(&xl, &self.data[l]);
            }
            if r % 2 == 0 {
                xr = M::op(&self.data[r - 1], &xr);
            }
            l = l / 2;
            r = (r - 1) / 2;
        }

        M::op(&xl, &xr)
    }
}

モノイド$(S, \bullet(op), id)$は次の様に定義できます。

pub struct MinMonoid;

impl Monoid for MinMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        std::cmp::min(*a, *b)
    }
    fn id() -> Self::S {
        std::usize::MAX
    }
}

pub struct MaxMonoid;

impl Monoid for MaxMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        std::cmp::max(*a, *b)
    }
    fn id() -> Self::S {
        0
    }
}

pub struct AddMonoid;

impl Monoid for AddMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        *a + *b
    }
    fn id() -> Self::S {
        0
    }
}

pub struct MulMonoid;

impl Monoid for MulMonoid {
    type S = usize;
    fn op(a: &Self::S, b: &Self::S) -> Self::S {
        *a * *b
    }
    fn id() -> Self::S {
        1
    }
}

  1. AOJ DSL_2_A Range Min Query (RMQ)【問題 / 解答例
  2. AOJ DSL_2_B Range Sum Query (RSQ)【問題 / 解答例
  3. AOJ DSL_2_E Range Add Query【問題 / 解答例
  4. AtCoder Beginner Contest 125 C - GCD on Blackboard 【問題 / 解答例】(+ 最大公約数)

アルゴリズム

二分探索 (Binary search, lower bound, upper bound)

  • めぐる式二分探索法によってlower bound / upper boundを実装すると次の様になります。
  • 配列の要素にOrd triatが実装されている必要があります。
  • isize型とusize型の間で足し算などの計算をする場合、適当に変換する必要があります。
pub trait BinarySearch<T> {
    fn lower_bound(&self, key: T) -> usize;
    fn upper_bound(&self, key: T) -> usize;
}

impl<T> BinarySearch<T> for [T]
where
    T: Ord,
{
    fn lower_bound(&self, key: T) -> usize {
        let mut ng = -1 as isize;
        let mut ok = self.len() as isize;
        while ok - ng > 1 {
            let mid = (ok + ng) / 2;
            if key <= self[mid as usize] {
                ok = mid;
            } else {
                ng = mid;
            }
        }
        ok as usize
    }

    fn upper_bound(&self, key: T) -> usize {
        let mut ng = -1 as isize;
        let mut ok = self.len() as isize;
        while ok - ng > 1 {
            let mid = (ok + ng) / 2;
            if key < self[mid as usize] {
                ok = mid;
            } else {
                ng = mid;
            }
        }
        ok as usize
    }
}

  1. AtCoder Beginner Contest 077 C - Snuke Festival【問題 / 解答例

(lower bound / upper bound以外)

  1. 競プロ典型 90 問 001 - Yokan Party(★4)【問題 / 解答例

再帰

  • フィボナッチ数を求めるサンプルです
fn fib(n: usize) -> usize {
    match n {
        0 => 0,
        1 => 1,
        _ => fib(n - 2) + fib(n - 1)
    }
}

fn main(){
    println!("{}", fib(10)); // 55
}

  1. AtCoder Beginner Contest 247 C - 1 2 1 3 1 2 1【問題 / 解答例

メモ化再帰

  • フィボナッチ数を求めるサンプルです。
  • Rustでは気軽にグローバル変数を扱えないので、メモをMapやVectorで持ち、必要に応じて再利用するようにすると実装できます。
use std::collections::HashMap;

fn fib(n: usize, memo: &mut HashMap<usize, u128>) -> u128 {
    if let Some(value) = memo.get(&n) {
        return *value;
    }

    let res = match n {
        0 => 0,
        1 => 1,
        _ => fib(n - 2, memo) + fib(n - 1, memo)
    };

    memo.insert(n, res);
    res
}

fn main(){
    let mut memo = HashMap::new();
    println!("{}", fib(100, &mut memo)); // 354224848179261915075
}
  1. AtCoder Beginner Contest 275 D - Yet Another Recursive Function 【問題 / 解答例

幅優先探索 (BFS, breadth first search)

  • データ構造の章で述べたように、先入れ先出し(FIFO)のQueueはstd::collections::VecDequeで実現できます。
  • 『AtCoder Beginner Contest 007 C - 幅優先探索』の解法サンプルです。
  • while letを使うと実装が楽になります。
use proconio::{
    input,
    marker::{Chars, Usize1},
};
use std::collections::VecDeque;

fn main() {
    input! {
        h: usize, w: usize,
        sy: Usize1, sx: Usize1,
        gy: Usize1, gx: Usize1,
        c: [Chars; h],
    }

    let mut que = VecDeque::new();
    que.push_back((sy, sx));
    let mut dist: Vec<Vec<Option<usize>>> = vec![vec![None; w]; h];
    dist[sy][sx] = Some(0);

    while let Some((y, x)) = que.pop_front() {
        if (y, x) == (gy, gx) {
            println!("{}", dist[y][x].unwrap());
            return;
        }

        let dy = vec![1, 0, -1, 0];
        let dx = vec![0, 1, 0, -1];

        for i in 0..4 {
            let dy = dy[i];
            let dx = dx[i];

            let next_y = y as isize + dy;
            let next_x = x as isize + dx;

            if 0 <= next_y
                && next_y < h as isize
                && 0 <= next_x
                && next_x < w as isize
                && dist[next_y as usize][next_x as usize].is_none()
                && c[next_y as usize][next_x as usize] == '.'
            {
                que.push_back((next_y as usize, next_x as usize));
                dist[next_y as usize][next_x as usize] = Some(dist[y][x].unwrap() + 1);
            }
        }
    }
}

  1. AtCoder Beginner Contest 007 C - 幅優先探索 【問題 / 解答例
  2. AtCoder Beginner Contest 272 D - Root M Leaper 【問題 / 解答例

ダイクストラ法 (Dijkstra's algorithm)

  • 次の様に実装できます。
  • 『AOJ 単一始点最短経路』の解法サンプルです。
  • 上で述べた様に、優先度付きキューはstd::collections::BinaryHeapで実装できます。
    • RustのBinaryHeapでは、標準で大きい順に要素が取り出されます。
    • State構造体を実装し、適切に順序関係(Ord, PartialEq)を入れることによって、小さい順に取り出せるようにします。
use std::{cmp::Ordering, collections::BinaryHeap};

const INF: usize = std::usize::MAX;

#[derive(Clone, PartialEq, Eq)]
pub struct State<T> {
    pub id: usize,
    pub priority: T,
}

impl<T> Ord for State<T>
where
    T: Ord,
{
    fn cmp(&self, other: &Self) -> Ordering {
        self.priority.cmp(&other.priority).reverse()
    }
}

impl<T> PartialOrd for State<T>
where
    T: Ord,
{
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

#[derive(Clone, PartialEq, Eq)]
pub struct Edge<T> {
    pub src: usize,
    pub dst: usize,
    pub weight: T,
}

#[derive(Clone)]
pub struct Graph<T> {
    pub edges: Vec<Vec<Edge<T>>>,
}

impl<T> Graph<T>
where
    T: Clone,
{
    pub fn new(n: usize) -> Self {
        Graph {
            edges: vec![vec![]; n],
        }
    }

    pub fn len(&self) -> usize {
        self.edges.len()
    }

    pub fn add_edge(&mut self, src: usize, dst: usize, weight: T) {
        self.edges[src].push(Edge { src, dst, weight });
    }
}

fn main() {
    let mut line = String::new();
    std::io::stdin().read_line(&mut line).unwrap();
    let word_list: Vec<&str> = line.split_whitespace().collect();
    let v_size: usize = word_list[0].parse::<usize>().unwrap();
    let e_size: usize = word_list[1].parse::<usize>().unwrap();
    let r: usize = word_list[2].parse::<usize>().unwrap();

    let mut std: Vec<(usize, usize, usize)> = Vec::new();
    for _ in 0..e_size {
        line = String::new();
        std::io::stdin().read_line(&mut line).unwrap();
        let word_list: Vec<usize> = {
            let word_list: Vec<&str> = line.split_whitespace().collect();
            word_list.into_iter().map(|s| s.parse::<usize>().unwrap()).collect()
        };
        std.push((word_list[0], word_list[1], word_list[2]));
    }

    let mut g = Graph::new(v_size);
    for &(u, v, c) in std.iter() {
        g.add_edge(u, v, c);
    }

    let mut dist = vec![INF; g.len()];
    dist[r] = 0;

    let mut que = BinaryHeap::new();
    que.push(State { id: r, priority: 0 });

    while let Some(State { id, priority }) = que.pop() {
        if priority > dist[id] {
            continue;
        }
        for e in &g.edges[id] {
            if dist[e.dst] > (dist[e.src]).checked_add(e.weight).unwrap_or(INF) {
                dist[e.dst] = dist[e.src] + e.weight;
                que.push(State {
                    id: e.dst,
                    priority: dist[e.dst],
                });
            }
        }
    }

    for i in 0..v_size {
        let ans = dist[i];
        if ans < INF {
            println!("{}", ans);
        } else {
            println!("INF");
        }
    }
}

  1. AOJ 単一始点最短経路【問題 / 解答例

最小全域木 (クラスカル法, Kruskal's algorithm)

  • Union-Findを用いて閉路の判定を行います。
  • 最小全域木の一つを構成する辺のVectorを返す実装です。
  • flatten()を使うと実装が楽です。
impl<T> Graph<T>
where
    T: Clone + Copy + Ord,
{
    pub fn kruskal(&self) -> Vec<Edge<T>> {
        let mut res = vec![];
        let mut uf = UnionFind::new(self.len());

        let mut all_edges = self
            .edges
            .iter()
            .cloned()
            .flatten()
            .collect::<Vec<Edge<T>>>();

        all_edges.sort_by_key(|e| e.weight);

        for e in all_edges {
            if !uf.is_same(e.src, e.dst) {
                uf.union(e.src, e.dst);
                res.push(e);
            }
        }

        res
    }
}

  1. AOJ GRL_2_A 最小全域木【問題 / 解答例

最大公約数/最小公倍数 (GCD/LCM)

  • ユークリッドの互除法を実装します。
  • 上で述べたstd::mem::swap関数を用いるとswapができます。
use std::mem;

pub fn gcd(mut a: usize, mut b: usize) -> usize {
    if a < b {
        mem::swap(&mut a, &mut b);
    }
    while b != 0 {
        let temp = a % b;
        a = b;
        b = temp;
    }
    a
}
/*
Rust 1.59.0からは分割代入(destructuring assignment)ができるようになったため
pub fn gcd(mut a: usize, mut b: usize) -> usize {
    if a < b {
        (a, b) = (b, a);
    }
    while b != 0 {
        (a, b) = (b, a % b);
    }
    a
}
の様にtempを挟まなくても動きます。swapも簡潔に書けます。各種コンテストサイトの言語アップデートが楽しみです。(2022/9)
*/

pub fn lcm(a: usize, b: usize) -> usize {
    a * b / gcd(a, b)
}
  • gcd()は、再帰を用いることによって、次のように書くこともできます
pub fn gcd(a: usize, b: usize) -> usize {
    if b == 0 {
        a
    } else {
        gcd(b, a % b)
    }
}

素数判定

  • $N$ に対して $\sqrt{N}$までの数が約数かどうか判定するアルゴリズムです。
  • take_whileを使うことで簡単に約数の候補を絞れます。
pub fn is_prime(n: usize) -> bool {
    if n == 0 || n == 1 {
        return false;
    }

    for i in (2..).take_while(|&x| x * x <= n) {
        if n % i == 0 {
            return false;
        }
    }

    true
}

  1. AtCoder Beginner Contest 084 D - 2017-like Number【問題 / 解答例

約数列挙

  • $\sqrt{N}$まで判定するアルゴリズムです。
  • 素数判定と同様にtake_whileが便利です。
pub fn gen_divisors(n: usize) -> Vec<usize> {
    let mut res = vec![];

    for i in (1..).take_while(|&x| x * x <= n) {
        if n % i == 0 {
            if i * i == n {
                res.push(i);
            } else {
                res.push(i);
                res.push(n / i);
            }
        }
    }
    res.sort();

    res
}

fn main() {
    println!("{:?}", gen_divisors(36)); // [1, 2, 3, 4, 6, 9, 12, 18, 36]
}

  1. AtCoder Beginner Contest 180 C - Cream puff【問題 / 解答例

素因数分解 (Integer factorization)

  • $n \in \mathbb{N},n = p_{1}^{e_{1}}p_{2}^{e_{2}}\dots p_{k}^{e_{k}}$と分解されるとき、$[(p_{1}: e_{1}), (p_{2}: e_{2}), \dots , (p_{k}: e_{k})]$なるmapを返す実装です。
  • mapのentryを使うと指数を簡単に記録できます。
  • BTreeMapで実装していますが、HashMapでも問題無いです。
use std::collections::BTreeMap;

pub fn integer_factorization(mut n: usize) -> BTreeMap<usize, usize> {
    let mut map = BTreeMap::new();

    let mut i = 2;
    while i * i <= n {
        while n % i == 0 {
            *map.entry(i).or_insert(0) += 1;
            n /= i;
        }
        i += 1;
    }

    if n != 0 && n != 1 {
        *map.entry(n).or_insert(0) += 1;
    }

    map
}
fn main() {
    println!("{:?}", integer_factorization(12)); // {2: 2, 3: 1}
    println!("{:?}", integer_factorization(100)); // {2: 2, 5: 2}
}
37
48
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
37
48