Rust snippets for competitive programing

libraries often used

https://docs.rs/proconio/latest/proconio/
https://doc.rust-lang.org/std/collections/struct.VecDeque.html
https://docs.rs/ordered-float/latest/ordered_float/struct.NotNan.html : float wrapper for Ord etc.

chmin

macro_rules! chmin {
    ($base:expr, $($cmps:expr),+ $(,)*) => {{
        let cmp_min = min!($($cmps),+);
        if $base > cmp_min {
            $base = cmp_min;
            true
        } else {
            false
        }
    }};
}

macro_rules! min {
    ($a:expr $(,)*) => {{
        $a
    }};
    ($a:expr, $b:expr $(,)*) => {{
        std::cmp::min($a, $b)
    }};
    ($a:expr, $($rest:expr),+ $(,)*) => {{
        std::cmp::min($a, min!($($rest),+))
    }};
}

chmax

macro_rules! chmax {
    ($base:expr, $($cmps:expr),+ $(,)*) => {{
        let cmp_max = max!($($cmps),+);
        if $base < cmp_max {
            $base = cmp_max;
            true
        } else {
            false
        }
    }};
}

macro_rules! max {
    ($a:expr $(,)*) => {{
        $a
    }};
    ($a:expr, $b:expr $(,)*) => {{
        std::cmp::max($a, $b)
    }};
    ($a:expr, $($rest:expr),+ $(,)*) => {{
        std::cmp::max($a, max!($($rest),+))
    }};
}

`[(usize, usize)]` to `Vec<HashSet>`

use std::collections::HashSet;
    
    let graph: Vec<HashSet<usize>> = {
        let mut tmp: Vec<HashSet<usize>> = vec![HashSet::new(); n];
        for path in paths.iter() {
            tmp[path.0 - 1].insert(path.1 - 1); // 0-index
        }
        tmp
    };

Heap

struct Heap {
    heap: Vec<usize>,
}

impl Heap {
    fn new() -> Self {
        Heap { heap: vec![] }
    }

    fn push(&mut self, x: usize) {
        self.heap.push(x);
        let mut i = self.heap.len() - 1;
        while i > 0 {
            let p = (i - 1) / 2;
            if self.heap[p] >= x {
                break;
            }
            self.heap[i] = self.heap[p];
            i = p;
        }
        self.heap[i] = x;
    }

    fn top(&mut self) -> Option<usize> {
        if !self.heap.is_empty() {
            return Some(self.heap[0]);
        };
        None
    }

    fn pop(&mut self) {
        if self.heap.is_empty() {
            return;
        };
        let x = self.heap.pop().unwrap();

        let mut i = 0;
        while (i * 2) + 1 < self.heap.len() {
            let mut child1 = i * 2 + 1;
            let child2 = i * 2 + 2;
            if child2 < self.heap.len() && self.heap[child2] > self.heap[child1] {
                child1 = child2;
            }
            if self.heap[child1] <= x {
                break;
            }
            self.heap[i] = self.heap[child1];
            i = child1;
        }
        self.heap[i] = x;
    }
}

dijkstra with `std::collections::BinaryHeap`

use std::cmp::Reverse;
use std::collections::BinaryHeap;
    // hogehoge //

    let mut graph = vec![vec![]; n];

    // fugafuga //
    let mut dist = vec![INF; n];
    let mut heap = BinaryHeap::new();
    heap.push((Reverse(0), 0));
    dist[0] = 0;

    while !heap.is_empty() {
        let (Reverse(d), v) = heap.pop().unwrap();
        if dist[v] != d {
            continue;
        };
        for &(to, cost) in &graph[v] {
            if chmin!(dist[to], dist[v] + cost) {
                heap.push((Reverse(dist[to]), to));
            }
        }
    }

u64::MAX

the Rust version of AtCoder is 1.42, u64::MAX hasn't implemented https://doc.rust-lang.org/std/primitive.u64.html#associatedconstant.MAX

const INF: u64 = 18446744073709551615;

input weighted direct graph

        let g = {
            let mut gg = vec![vec![]; n];
            for _ in 0..m {
                input! {
                    u: Usize1,
                    v: Usize1,
                    c: usize,
                }
                gg[u].push((v, c));
            }
            gg
        };

Union-Find

use std::mem::swap;

struct UnionFind {
    par: Vec<usize>,
    siz: Vec<usize>,
}

impl UnionFind {
    fn new(n: usize) -> Self {
        UnionFind {
            par:(0..n).collect(),
            siz: vec![1; n],
        }
    }

    fn root(&self, x: usize) -> usize {
        if self.par[x] == x {
            return x
        }
        self.root(self.par[x])
    }

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

    fn unite(&mut self, x: usize, y: usize) -> bool {
        let mut x_root = self.root(x);
        let mut y_root = self.root(y);

        if x_root == y_root {
            return false
        }
        if self.siz[x_root] < self.siz[y_root] {
            swap(&mut x_root, &mut y_root);
        }
        self.par[y_root]= x_root;
        self.siz[x_root] += self.siz[y_root];
        true
    }

    fn size(&self, x: usize) -> usize {
        self.siz[self.root(x)]
    }
    
    fn is_connected(&self) -> bool {
        for i in 1..self.par.len() {
            if !self.is_same(0, i) {
                return false;
            }
        }
        true
    }    
}

stdin n, m without proconio

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    let mut s = String::new();
    std::io::stdin().read_line(&mut s).ok();
    s.trim()
        .split_whitespace()
        .map(|e| e.parse().ok().unwrap())
        .collect()
}

fn main() {
    let (n, m) = {
        let v: Vec<usize> = read_vec();
        (v[0], v[1])
    };
}

more details:
https://codeforces.com/blog/entry/96067
https://qiita.com/penguinshunya/items/cd96803b74635aebefd6

// 5 5
// .....
// .#.#.
// .###.
// .#.#.
// .....

use std::io::BufRead;

fn read_vec<T: std::str::FromStr>() -> Vec<T> {
    let mut s = String::new();
    std::io::stdin().read_line(&mut s).ok();
    s.trim()
        .split_whitespace()
        .map(|e| e.parse().ok().unwrap())
        .collect()
}

fn read_chars() -> Vec<char> {
    let mut s = String::new();
    std::io::stdin().read_line(&mut s).ok();
    trim_newline(&mut s);
    s.chars().collect()
}

fn trim_newline(s: &mut String) {
    if s.ends_with('\n') {
        s.pop();
        if s.ends_with('\r') {
            s.pop();
        }
    }
}

fn main() {
    let (r, c) = {
        let v: Vec<usize> = read_vec();
        (v[0], v[1])
    };
    let s: Vec<Vec<char>> = {
        let mut tmp: Vec<Vec<char>> = vec![vec![]; r];
        for i in 0..r {
            let mut v: Vec<char> = read_chars();
            tmp[i].append(&mut v);
        }
        tmp
    };
    assert_eq!("[['.', '.', '.', '.', '.'], ['.', '#', '.', '#', '.'], ['.', '#', '#', '#', '.'], ['.', '#', '.', '#', '.'], ['.', '.', '.', '.', '.']]", format!("{:?}", s))
}

bit brute force

ex. return the list of proper parenthesis

    pub fn generate_parenthesis(n: i32) -> Vec<String> {
        let mut ans = vec![];
        for bit in 0..(1 << n*2) {
            let s = (0..n*2)
                .map(|x| if bit & (1 << x) != 0 { '(' } else { ')' })
                .collect::<String>();
            if Self::ok(s.clone()) {
                ans.push(s);
            }
        }
        ans
    }

in Python:

class Solution:
    def subsets(self, nums: List[int]) -> List[List[int]]:
        n = len(nums)
        ans = []
        for bit in range(1 << n):
            v = []
            for j in range(n):
                if bit & (1 << j) != 0:
                    v.append(nums[j])
            ans.append(v)
        return ans

stdout with buffering

    use std::io::{stdout, BufWriter, Write};

    let mut writer = BufWriter::new(stdout());
    writeln!(writer, "hoge");
    writeln!(writer, "fuga");
    writer.flush().unwrap();
    // hoge
    // fuga

bit SUM and XOR

(a XOR b) + 2 * (a AND b) = a + b
// more details:
// https://stackoverflow.com/questions/36477623/given-an-xor-and-sum-of-two-numbers-how-to-find-the-number-of-pairs-that-satisf
// https://atcoder.jp/contests/abc238/tasks/abc238_d

Ford Fulkerson

pub mod ford_fullkerson {
    use std::{
        cmp::min,
        collections::{HashMap, HashSet},
    };
    pub struct Graph {
        edges: Vec<HashMap<usize, usize>>,
    }

    const INF: usize = 18446744073709551615;

    impl Graph {
        pub fn new(n: usize) -> Self {
            Graph {
                edges: vec![HashMap::new(); n],
            }
        }
        pub fn add_edges(&mut self, from: usize, to: usize, capacity: usize) {
            *self.edges[from].entry(to).or_insert(0) += capacity;
            self.edges[to].entry(from).or_insert(0);
        }
        fn run_flow(&mut self, from: usize, to: usize, flow: usize) {
            *self.edges[from].entry(to).or_insert(0) -= flow;
            *self.edges[to].entry(from).or_insert(0) += flow;
        }
        pub fn max_flow(&mut self, start: usize, goal: usize) -> usize {
            let mut flow = 0;
            loop {
                let f = self.dfs(start, goal, INF, &mut HashSet::new());
                if f == 0 {
                    break;
                } else {
                    flow += f;
                }
            }
            flow
        }
        pub fn dfs(
            &mut self,
            now: usize,
            goal: usize,
            flow: usize,
            visited: &mut HashSet<usize>,
        ) -> usize {
            if now == goal {
                return flow;
            }
            visited.insert(now);
            for (dest, capacity) in self.edges[now].clone() {
                if !visited.contains(&dest) && capacity > 0 {
                    let f = self.dfs(dest, goal, min(flow, capacity), visited);
                    if f > 0 {
                        self.run_flow(now, dest, f);
                        return f;
                    }
                }
            }
            0
        }
    }
}

Ford Fulkerson (allow duplicate edge)

pub mod ford_fullkerson {
    use std::cmp;
    const INF: u64 = 18446744073709551615;

    #[derive(Debug, Copy, Clone)]
    struct Edge {
        from: usize,
        to: usize,
        cap: u64,
        index: usize,
        rev_index: usize,
    }

    pub struct Graph {
        list: Vec<Vec<Edge>>,
    }

    impl Graph {
        pub fn new(v_n: usize) -> Self {
            Self {
                list: vec![Vec::new(); v_n],
            }
        }

        fn size(&self) -> usize {
            self.list.len()
        }

        fn edge<'a>(&'a mut self, e: Edge) -> &'a mut Edge {
            &mut self.list[e.from][e.index]
        }

        fn redge<'a>(&'a mut self, e: Edge) -> &'a mut Edge {
            &mut self.list[e.to][e.rev_index]
        }

        fn run_flow(&mut self, e: Edge, f: u64) {
            self.edge(e).cap -= f;
            self.redge(e).cap += f;
        }

        pub fn addedge(&mut self, from: usize, to: usize, cap: u64) {
            let last_from_index = self.list[from].len();
            let last_to_index = self.list[to].len();

            self.list[from].push(Edge {
                from,
                to,
                cap,
                index: last_from_index,
                rev_index: last_to_index,
            });
            self.list[to].push(Edge {
                from: to,
                to: from,
                cap: 0,
                index: last_to_index,
                rev_index: last_from_index,
            });
        }

        fn dfs(&mut self, seen: &mut Vec<bool>, v: usize, t: usize, f: u64) -> Option<u64> {
            if v == t {
                return Some(f);
            }

            seen[v] = true;
            let from_v_edges = self.list[v].clone();

            for e in from_v_edges.into_iter() {
                if seen[e.to] {
                    continue;
                }

                if self.edge(e).cap == 0 {
                    continue;
                }
                let e_cap = self.edge(e).cap;
                match self.dfs(seen, e.to, t, cmp::min(f, e_cap)) {
                    None => {
                        continue;
                    }
                    Some(flow) => {
                        self.run_flow(e, flow);
                        return Some(flow);
                    }
                }
            }

            None
        }
        pub fn max_flow(&mut self, s: usize, t: usize) -> u64 {
            let mut max_flow = 0_u64;
            loop {
                let mut seen = vec![false; self.size()];
                match self.dfs(&mut seen, s, t, INF) {
                    None => {
                        break;
                    }
                    Some(flow) => {
                        max_flow += flow;
                    }
                }
            }
            max_flow
        }
    }
}

Floyd-Warshal

    // Floyd-Warshal
    // Init
    let mut dp = vec![vec![INF; n]; n];
    for _ in 0..m {
        input! {
            a: usize, b: usize, w: usize
        }
        dp[a][b] = w;
        dp[b][a] = w;
    }
    for i in 0..n {
        dp[i][i] = 0;
    }

    // DP
    for k in 0..n {
        for i in 0..n {
            for j in 0..n {
                chmin!(dp[i][j], dp[i][k] + dp[k][j]);
            }
        }
    }

逆元nCr

const MOD: u64 = 1_000_000_007;

fn inv(mut a: u64) -> u64 {
    let mut t = 1;
    let mut n = MOD - 2;
    while n > 0 {
        if n & 1 == 1 {
            t = t * a % MOD;
        }
        a = a * a % MOD;
        n >>= 1;
    }
    t
}

fn main() {
    input!(n: usize);
    let mut fact = vec![1; n + 1];
    for i in 1..=n {
        fact[i] = i as u64 * fact[i - 1] % MOD;
    }
    let mut ifact = vec![1; n + 1];
    ifact[n] = inv(fact[n]);
    for i in (1..n).rev() {
        ifact[i] = (i + 1) as u64 * ifact[i + 1] % MOD;
    }
    let combination = |n: usize, r: usize| -> u64 { fact[n] * ifact[r] % MOD * ifact[n - r] % MOD }
    
    // hogehoge
}

fenwick

mod fenwick {
    #[derive(Debug)]
    pub struct FenwickTree {
        tree: Vec<i64>,
    }

    impl FenwickTree {
        pub fn new(n: usize) -> Self {
            Self {
                tree: vec![0; n + 1],
            }
        }

        pub fn add(&mut self, mut i: usize, x: i64) {
            i += 1;
            while i < self.tree.len() {
                self.tree[i] += x;
                i += i & i.wrapping_neg();
            }
        }

        pub fn sum(&self, mut i: usize) -> i64 {
            i += 1;
            let mut result = 0;
            while i > 0 {
                result += self.tree[i];
                i ^= i & i.wrapping_neg();
            }
            result
        }
    }
}

min BinaryHeap from vec

    let mut b: BinaryHeap<_> = v.into_iter().map(Reverse).collect();

create sum vec with scan

    let mut s = a
        .iter()
        .scan(0, |acc, x| {
            *acc += *x;
            Some(*acc)
        })
        .collect::<Vec<_>>();

座標圧縮

use proconio::input;
use proconio::marker::Chars;
use std::collections::HashSet;

fn main() {
    input! {
        h: usize, w: usize,
        mut a: [Chars; h]
    }
    let mut vx = HashSet::new();
    let mut vy = HashSet::new();
    for i in 0..h {
        if !a[i].contains(&'#') {
            vy.insert(i);
        }
    }
    for i in 0..w {
        let mut f = true;
        for j in (i..w * h).step_by(w) {
            if *a.iter().flatten().collect::<Vec<_>>()[j] == '#' {
                f = false;
            }
        }
        if f {
            vx.insert(i);
        }
    }
    for i in 0..h {
        if vy.contains(&i) {
            continue;
        }
        for j in 0..w {
            if vx.contains(&j) {
                continue;
            }
            print!("{}", a[i][j]);
        }
        println!()
    }
}

天井関数ceil 切り上げ

gcd

fn gcd(first: usize, second: usize) -> usize {
    let mut max = first;
    let mut min = second;
    if min > max {
        let val = max;
        max = min;
        min = val;
    }

    loop {
        let res = max % min;
        if res == 0 {
            return min;
        }

        max = min;
        min = res;
    }
}

kubosuke/rust-cp-snippets.md

libraries often used

chmin

chmax

[(usize, usize)] to Vec<HashSet>

Heap

dijkstra with std::collections::BinaryHeap

u64::MAX

input weighted direct graph

Union-Find

stdin n, m without proconio

bit brute force

stdout with buffering

bit SUM and XOR

Ford Fulkerson

Ford Fulkerson (allow duplicate edge)

Floyd-Warshal

逆元nCr

fenwick

min BinaryHeap from vec

create sum vec with scan

座標圧縮

天井関数ceil 切り上げ

gcd

`[(usize, usize)]` to `Vec<HashSet>`

dijkstra with `std::collections::BinaryHeap`