AtCoder Rust Tips集など

使えるRustおよびライブラリのバージョンに関して

に記載されている感じ。

2024/2/17時点で、本体のバージョンは1.70.0

使えるクレート

ac-library-rs@=0.1.1
once_cell@=1.18.0
static_assertions@=1.1.0
varisat@=0.2.2
memoise@=0.3.2
argio@=0.2.0
bitvec@=1.0.1
counter@=0.5.7
hashbag@=0.1.11
pathfinding@=4.3.0
recur-fn@=2.2.0
indexing@=0.4.1
amplify@=3.14.2
amplify_derive@=2.11.3
amplify_num@=0.4.1
easy-ext@=1.0.1
multimap@=0.9.0
btreemultimap@=0.1.1
bstr@=1.6.0
az@=1.2.1
glidesort@=0.1.2
tap@=1.0.1
omniswap@=0.1.0
multiversion@=0.7.2
num@=0.4.1
num-bigint@=0.4.3
num-complex@=0.4.3
num-integer@=0.1.45
num-iter@=0.1.43
num-rational@=0.4.1
num-traits@=0.2.15
num-derive@=0.4.0
ndarray@=0.15.6
nalgebra@=0.32.3
alga@=0.9.3
libm@=0.2.7
rand@=0.8.5
getrandom@=0.2.10
rand_chacha@=0.3.1
rand_core@=0.6.4
rand_hc@=0.3.2
rand_pcg@=0.3.1
rand_distr@=0.4.3
petgraph@=0.6.3
indexmap@=2.0.0
regex@=1.9.1
lazy_static@=1.4.0
ordered-float@=3.7.0
ascii@=1.1.0
permutohedron@=0.2.4
superslice@=1.0.0
itertools@=0.11.0
itertools-num@=0.1.3
maplit@=1.0.2
either@=1.8.1
im-rc@=15.1.0
fixedbitset@=0.4.2
bitset-fixed@=0.1.0
proconio@=0.4.5
text_io@=0.1.12
rustc-hash@=1.1.0
smallvec@=1.11.0

3rdパーティーライブラリ無しで頑張るのも一興だが、標準入力の読み込みが大変なのでproconio:

などは便利。

その他で(個人的に)ときどき使う、あるいは使いそうなやつだと

https://crates.io/crates/superslice

など。

7e+8

AtCoder + Rustに関するまとめ記事など
（一部Rust関係ないものも含む）

（便利そうなものがあれば適宜追加）

7e+8

UnionFind

petgraph:

で使うことができる。

使ってみた例:

実装(ABC214 D)

use petgraph::unionfind::UnionFind;
use proconio::input;

fn main() {
    input! {
        n: usize,
        e: [(usize, usize, i128); n - 1],
    }
    let mut e = e
        .iter()
        .map(|(x, y, w)| (*w, *x - 1, *y - 1))
        .collect::<Vec<_>>();
    e.sort();

    let mut uf = UnionFind::<usize>::new(n);
    let mut tree_size_tmp = vec![1i128; n];
    let mut sum = 0i128;

    for (w, x, y) in e {
        let i = uf.find(x);
        let j = uf.find(y);
        sum += w * tree_size_tmp[i] * tree_size_tmp[j];

        uf.union(x, y);
        tree_size_tmp[uf.find(x)] = tree_size_tmp[i] + tree_size_tmp[j];
    }

    println!("{sum}");
}

なお、petgraphを使わずに自分でUnionFindを実装した場合は↓のような感じ

実装例(ABC214 D petgraphなし)

use proconio::input;

type K = usize;

struct UnionFind {
    parent: Vec<K>,
    rank: Vec<usize>,
}

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

    fn root(&mut self, x: K) -> K {
        if self.parent[x] == x {
            x
        } else {
            let r = self.root(self.parent[x]);
            self.parent[x] = r;
            r
        }
    } 

    fn is_same(&mut self, x: K, y: K) -> bool {
        self.root(x) == self.root(y)
    }

    fn merge(&mut self, x: K, y: K) -> bool {
        let x = self.root(x);
        let y = self.root(y);
        if x == y {
            return false;
        }

        if self.rank[x] < self.rank[y] {
            self.parent[x] = y;
        } else if self.rank[x] > self.rank[y] {
            self.parent[y] = x;
        } else {
            self.rank[x] += 1;
            self.parent[y] = x;
        }

        true
    }
}

fn main() {
    input! {
        n: usize,
        e: [(usize, usize, i128); n - 1],
    }
    let mut e = e
        .iter()
        .map(|(x, y, w)| (*w, *x - 1, *y - 1))
        .collect::<Vec<_>>();
    e.sort();

    let mut uf = UnionFind::new(n);
    let mut tree_size_tmp = vec![1i128; n];
    let mut sum = 0i128;

    for (w, x, y) in e {
        let i = uf.root(x);
        let j = uf.root(y);
        sum += w * tree_size_tmp[i] * tree_size_tmp[j];

        uf.merge(x, y);
        tree_size_tmp[uf.root(x)] = tree_size_tmp[i] + tree_size_tmp[j];
    }

    println!("{sum}");
}

↑に関して、

などを参考にしている

(追記)

UnionFindの講座用問題:

があったのでその例を追加

(ATC001 b)

// petgraph使用
use petgraph::unionfind::UnionFind;
use proconio::input;

fn main() {
    input! {
        n: usize,
        q: usize,
    }

    let mut uf = UnionFind::<usize>::new(n + 1);
    let mut answers = Vec::new();

    for _ in 0..q {
        input! {
            p: i8,
            a: usize,
            b: usize,
        }

        if p == 0 {
            uf.union(a, b);
        } else {
            if uf.equiv(a, b) {
                answers.push("Yes");
            } else {
                answers.push("No");
            }
        }
    }

    for ans in &answers {
        println!("{ans}");
    }
}

// petgraph不使用
use proconio::input;

type K = usize;

struct UnionFind {
    parent: Vec<K>,
    rank: Vec<usize>,
}

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

    fn root(&mut self, x: K) -> K {
        if self.parent[x] == x {
            x
        } else {
            let r = self.root(self.parent[x]);
            self.parent[x] = r;
            r
        }
    }

    fn is_same(&mut self, x: K, y: K) -> bool {
        self.root(x) == self.root(y)
    }

    fn merge(&mut self, x: K, y: K) -> bool {
        let x = self.root(x);
        let y = self.root(y);
        if x == y {
            return false;
        }

        if self.rank[x] < self.rank[y] {
            self.parent[x] = y;
        } else if self.rank[x] > self.rank[y] {
            self.parent[y] = x;
        } else {
            self.rank[x] += 1;
            self.parent[y] = x;
        }

        true
    }
}

fn main() {
    input! {
        n: usize,
        q: usize,
    }

    let mut uf = UnionFind::new(n + 1);
    let mut answers = Vec::new();

    for _ in 0..q {
        input! {
            p: i8,
            a: usize,
            b: usize,
        }

        if p == 0 {
            uf.merge(a, b);
        } else {
            if uf.is_same(a, b) {
                answers.push("Yes");
            } else {
                answers.push("No");
            }
        }
    }

    for ans in &answers {
        println!("{ans}");
    }
}

7e+8

Dijkstra法

例題: ABC214 C

petgraphを使わずにやった場合:

実装(ABC214 C)

use proconio::input;
use std::cmp::Reverse;

type K = i64;
const INF: K = 1i64 << 61;
struct Edge {
    from: usize,
    to: usize,
    weight: K,
}

fn dijkstra(edges: &Vec<Edge>, node_count: usize, node_from: usize) -> Vec<K> {
    let mut edges_from_nodes = vec![Vec::<(usize, K)>::new(); node_count];
    for Edge { from, to, weight } in edges {
        edges_from_nodes[*from].push((*to, *weight));
    }

    #[derive(Ord, PartialOrd, Eq, PartialEq)]
    struct Item(Reverse<K>, usize);

    let mut heap = std::collections::BinaryHeap::<Item>::new();
    heap.push(Item(Reverse(0), node_from));

    let mut shortest = vec![INF; node_count];

    while let Some(x) = heap.pop() {
        let Item(Reverse(score), node) = x;
        if score < shortest[node] {
            shortest[node] = score;
            for (to, w) in &edges_from_nodes[node] {
                heap.push(Item(Reverse(score + *w), *to));
            }
        }
    }

    shortest
}

fn main() {
    input! {
        n: usize,
        s: [i64; n],
        t: [i64; n],
    }

    let mut edges = Vec::<Edge>::new();
    for (i, &w) in s.iter().enumerate() {
        let i = i + 1;
        let j = i % n + 1;
        edges.push(Edge {
            from: i,
            to: j,
            weight: w,
        });
    }

    for (i, &w) in t.iter().enumerate() {
        let i = i + 1;
        edges.push(Edge {
            from: 0,
            to: i,
            weight: w,
        });
    }

    let v = dijkstra(&edges, n + 1, 0);

    for x in &v[1..] {
        println!("{x}");
    }
}

↑参考:

petgraphを使った場合:

実装(ABC214 C petgraph使用)

use petgraph::{algo::dijkstra, graph::DiGraph};
use proconio::input;

fn main() {
    input! {
        n: usize,
        s: [i64; n],
        t: [i64; n],
    }

    let mut g = DiGraph::<usize, i64>::new();

    let nodes = {
        let mut v = Vec::new();
        for i in 0..(n + 1) {
            let node = g.add_node(i);
            v.push(node);
        }
        v
    };

    for (i, &w) in s.iter().enumerate() {
        let i = i + 1;
        let j = (i % n) + 1;
        // e.push((nodes[i], nodes[j], w));
        g.add_edge(nodes[i], nodes[j], w);
    }

    for (i, &w) in t.iter().enumerate() {
        let j = i + 1;
        g.add_edge(nodes[0], nodes[j], w);
    }

    let answers = dijkstra(&g, nodes[0], None, |e| *e.weight());

    for key in &nodes[1..] {
        let ans = answers.get(key).unwrap();
        println!("{ans}");
    }
}

↑参考:

（petgraphのドキュメントが少ない+使用例が意外と見つからないなどで意外に苦労した）

7e+8

素因数分解

一例としては↓のような感じ

素因数分解実装例

(素数の値, 素数の個数)のVectorを返す関数の例:

type K = i64;
fn prime_factorization(n: K) -> Vec<(K, K)> {
    let mut n = n;
    let mut x = 2;

    let mut v = Vec::new();

    while x * x <= n {
        if n % x != 0 {
            x += 1;
            continue;
        }

        let mut cnt = 0;
        while n % x == 0 {
            cnt += 1;
            n /= x;
        }

        v.push((x, cnt));
    }

    if n != 1 {
        v.push((n, 1));
    }

    v
}

ARC067 Cでの実装例:

実装例(arc067 c)

use proconio::input;

const MOD: i64 = 1e9 as i64 + 7;

type K = i64;
fn prime_factorization(n: K) -> Vec<(K, K)> {
    let mut n = n;
    let mut x = 2;

    let mut v = Vec::new();

    while x * x <= n {
        if n % x != 0 {
            x += 1;
            continue;
        }

        let mut cnt = 0;
        while n % x == 0 {
            cnt += 1;
            n /= x;
        }

        v.push((x, cnt));
    }

    if n != 1 {
        v.push((n, 1));
    }

    v
}

fn main() {
    input! {
        n: i64,
    }

    if n == 1 {
        println!("1");
        return;
    }

    let mut m = std::collections::BTreeMap::new();
    for x in 2..(n + 1) {
        let v = prime_factorization(x);

        for (p, cnt) in v {
            if let Some(&old) = m.get(&p) {
                m.insert(p, old + cnt);
            } else {
                m.insert(p, cnt);
            }
        }
    }

    let mut ans = 1;

    for (_, cnt) in m {
        ans *= (cnt + 1) % MOD;
        ans %= MOD;
    }

    println!("{ans}");
}