競技プログラミングを行う上で使用頻度の高いアルゴリズムやデータ構造を、Rust(プログラミング言語)で実装するためのノートです。各アルゴリズムの実装と、それを用いる問題や解答例を掲載しています。
目次
関数/メソッド
データ構造
- 動的配列 (vector)
- スタック (stack)
- キュー (queue)
- 優先度付きキュー (priority queue)
- Set
- Map
- Union-Find木 (素集合データ構造)
- Mod Int (Mint, ModInt)
- Segment Tree (セグメント木)
アルゴリズム
グラフ関連
- 幅優先探索 (BFS, breadth first search)
- ダイクストラ法 (Dijkstra's algorithm)
- 最小全域木 (クラスカル法, Kruskal's algorithm)
整数論関連
関数/メソッド
最小値/最大値 (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
}
例
- AtCoder Beginner Contest 052 A - Two Rectangles【問題 / 解答例】
- 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
}
例
スワップ (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
}
例
ソート (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]
}
例
sort_by()
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]
}
例
また、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]
}
例
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)
- 標準のstd::collections::BinaryHeapで実現できます。
- 要素が大きい順に取り出される点に注意してください。
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
}
例
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
}
例
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)
}
}
例
- ある要素が属しているグループに含まれるノードの数を調べることができるように拡張すると、次のようになります。(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]
}
}
}
例
- AtCoder Beginner Contest 120 D - Decayed Bridges【問題 / 解答例】
- 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
}
例
- AtCoder Beginner Contest 034 C - 経路【問題 / 解答例】
- AtCoder Regular Contest 145 C - Split and Maximize 【問題 / 解答例】
Segment Tree (セグメント木)
-
RustではACL(AtCoder Library)が使えないので各自で実装をします。
-
モノイド$(S, \bullet(op), id)(S:$ 集合、$\bullet: S$上の二項演算、$id: S$の$\bullet$に関する単位元$)$を用いて、数列$\lbrace a_i \rbrace_{i=1}^n$に対して、次の1~3ができます。
- update: $a_{i}$ をxに更新
- get: $a_{i}$ の取得
- 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
}
}
例
- AOJ DSL_2_A Range Min Query (RMQ)【問題 / 解答例】
- AOJ DSL_2_B Range Sum Query (RSQ)【問題 / 解答例】
- AOJ DSL_2_E Range Add Query【問題 / 解答例】
- 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
}
}
例
例 (lower bound / upper bound以外)
再帰
- フィボナッチ数を求めるサンプルです
fn fib(n: usize) -> usize {
match n {
0 => 0,
1 => 1,
_ => fib(n - 2) + fib(n - 1)
}
}
fn main(){
println!("{}", fib(10)); // 55
}
例
メモ化再帰
- フィボナッチ数を求めるサンプルです。
- 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
}
幅優先探索 (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);
}
}
}
}
例
- AtCoder Beginner Contest 007 C - 幅優先探索 【問題 / 解答例】
- 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");
}
}
}
例
最小全域木 (クラスカル法, 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
}
}
例
最大公約数/最小公倍数 (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
}
例
約数列挙
- $\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]
}
例
素因数分解 (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}
}