atcoder rustチートシート

コンテストのときに自分が見る用。そのままだと動かない
ACLをふんだんに使用

所有権

for x in a {}          // aを消費。以後aを使えない
for &x in &a {}        // aを残して値を取り出す
for x in &mut a { *x += 1; }

fn calc(a: &[usize]) {}         // 読むだけ
fn update(a: &mut [usize]) {}   // 書き換える

iterator

a.iter().enumerate()               // (添字, &値)
a.iter().zip(b.iter())             // 2配列を同時に見る
a.iter().any(|&x| x > 0)           // 1つでも
a.iter().all(|&x| x > 0)           // 全て
a.iter().filter(|&&x| x > 0).count()
a.iter().sum::<i64>()
a.iter().copied().max().unwrap() // 最大値を取り出す

maxはiteratorのメソッドで、Vecには使えない

ソート・二分探索

a.sort_unstable();                         // 昇順
a.sort_unstable_by(|a, b| b.cmp(a));       // 降順
a.sort_unstable_by_key(|x| x.0);           // キー指定
a.dedup();                                 // 連続する重複を削除

// Pythonのbisect_left / bisect_right相当
let lower = a.partition_point(|&x| x < target);
let upper = a.partition_point(|&x| x <= target);

match a.binary_search(&target) {
    Ok(i) => {}   // 存在
    Err(i) => {}  // 挿入位置
}

partition_pointもbinary_searchも、事前にソートさえしておけばインポートせずに使える

HashMap/HashSet

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

let mut cnt = HashMap::new();
*cnt.entry(x).or_insert(0) += 1;

let value = cnt.get(&x).copied().unwrap_or(0);
let exists = cnt.contains_key(&x);

set.insert(x);
set.contains(&x);
set.remove(&x);

a.intersection(&b);
a.union(&b);
a.difference(&b);

position

やりたいこと	書き方
値と一致する位置	.position(|&x| x == target)
条件を満たす位置	.position(|&x| x >= 10)
構造体・Vecを比較	.position(|x| x == &target)
後ろから探す	.rposition(|&x| x == target)
絶対に存在する	.position(...).unwrap()
なければ既定値	.position(...).unwrap_or(default)

abs

// 符号付き整数1つの絶対値
let ans = x.abs();

// usize同士/i64同士の差
let ans = a.abs_diff(b);

// MINまで含めて絶対値を安全に扱う
let ans = x.unsigned_abs();

クロージャ

let f = |x: i32| x * 2;      // 引数あり
let g = || 10;               // 引数なし
let h = |a: i32, b: i32| a + b;
println!("{}", f(3));         // 6
// 外側の変数も使えるのが普通の関数との違い

文字列

use proconio::marker::{Chars, Bytes, Usize1};

input! {
    s: Chars,                  // Vec<char>
    t: Bytes,                  // Vec<u8>
    edges: [(Usize1, Usize1); m], // 1-indexed入力を0-indexedに
}

// ASCII
let s = s.as_bytes();
s[i] == b'a';
s.windows(2).any(|w| w == b"ab");

// Unicode文字
let s: Vec<char> = s.chars().collect();
s[i] == 'あ';

// 部分文字列
s.contains("abc");
s.starts_with("abc");
s.ends_with("abc");

BFS

let mut queue = VecDeque::new();
queue.push_back(start);

while let Some(v) = queue.pop_front() {
    // 次の頂点をqueueに入れる
}

// 詳細
use std::collections::VecDeque;

fn bfs(graph: &[Vec<usize>], start: usize) -> Vec<isize> {
    let n = graph.len();

    // -1は未到達
    let mut dist = vec![-1; n];
    let mut queue = VecDeque::new();

    dist[start] = 0;
    queue.push_back(start);

    while let Some(v) = queue.pop_front() {
        for &next in &graph[v] {
            if dist[next] != -1 {
                continue;
            }

            dist[next] = dist[v] + 1;
            queue.push_back(next);
        }
    }

    dist
}

bit全探索

for bit in 0..(1_usize << n) {
    for i in 0..n {
        if bit >> i & 1 == 1 {
            // iを選ぶ
        }
    }
}
bit.count_ones();      // 立っているbit数
bit.trailing_zeros();  // 最下位の1の位置

累積和

let mut sum = vec![0_i64; n + 1];

for i in 0..n {
    sum[i + 1] = sum[i] + a[i];
}

let range_sum = sum[r] - sum[l]; // a[l..r]

組み合わせ

欲しいもの	メソッド	例	個数
順番を無視して r 個選ぶ	.combinations(r)	{1,2} と {2,1} は同じ	({}_nC_r)
順番を区別して r 個選ぶ	.permutations(r)	[1,2] と [2,1] は別	({}_nP_r)
2個・3個の組をタプルで取得	.tuple_combinations()	(1,2), (1,3)	({}_nC_r)
2集合から1個ずつ選ぶ	.cartesian_product()	(a[0], b[0]) など	n × m
すべての部分集合	.powerset()	空集合も含む	(2^n)

// 異なる2要素を選ぶ
for (&x, &y) in a.iter().tuple_combinations() {}

// r個選ぶ
for v in a.iter().combinations(r) {}

// 並べる順番も重要
for v in a.iter().permutations(r) {}

// aとbから1個ずつ選ぶ
for (&x, &y) in a.iter().cartesian_product(b.iter()) {}

// 全部分集合を試す
for subset in a.iter().powerset() {}

use permutohedron::LexicalPermutation;
let mut p = vec![1, 2, 3];
p.next_permutation();         // 辞書順で次へ：[1, 3, 2]
let ok = p.next_permutation(); // 次があれば true、なければ false

ダイクストラ

let mut heap = BinaryHeap::new();
heap.push(Reverse((0_i64, start)));

while let Some(Reverse((d, v))) = heap.pop() {
    if d != dist[v] {
        continue;
    }
    // 距離を更新してheapに入れる
}

// 詳細
use std::cmp::Reverse;
use std::collections::BinaryHeap;

fn dijkstra(graph: &[Vec<(usize, i64)>], start: usize) -> Vec<i64> {
    let n = graph.len();
    const INF: i64 = 1_i64 << 60;

    let mut dist = vec![INF; n];
    let mut heap = BinaryHeap::new();

    dist[start] = 0;

    // Reverseを使い、距離が小さいものから取り出す
    heap.push(Reverse((0_i64, start)));

    while let Some(Reverse((current_dist, v))) = heap.pop() {
        // 古い情報なら無視
        if current_dist != dist[v] {
            continue;
        }

        for &(next, cost) in &graph[v] {
            let next_dist = current_dist + cost;

            if next_dist < dist[next] {
                dist[next] = next_dist;
                heap.push(Reverse((next_dist, next)));
            }
        }
    }

    dist
}

Union-Find

use ac_library::Dsu;

let mut uf = Dsu::new(5);
uf.merge(0,1);
uf.same(0,2): bool;
uf.size(0); // 0が属するグループの要素数
uf.leader(0); // 0が属するグループの代表元
uf.groups(); // 全グループ

セグメント木

配列の値を一点ずつ更新しながら、区間の合計・最小値・最大値などを何度も求める場合

use ac_library::segtree::{Segtree, Max, Min, Additive};

let a = !vec[0;n]
let mut seg = Segtree::<>::from(a);

// Segtree::<Max<i64>>   // 区間最大
// Segtree::<Min<i64>>   // 区間最小
// Segtree::<Additive<i64>> // 区間和

seg.set(2,10); // a[2]を10に更新
seg.prod(1..4); // a[1..4]の区間〇〇を取得
seg.get(2); // 一点取得
// 二分探索 条件sum(l..r) <= xを満たす最小/最大の相方を探す
seg.max_right(l, |&sum| sum <= x);
seg.min_left(r, |&sum| sum <= x);

二分探索のユースケース例：「l から足して、合計が x 以下になる最大の右端」を探す

遅延セグ木

配列の区間全体をまとめて更新しながら、区間の合計・最小値・最大値などを何度も求める場合

use ac_library::{LazySegtree, MapMonoid, Max};

struct RangeAddMax;

impl MapMonoid for RangeAddMax {
    type M = Max<i64>;
    type F = i64; // 遅延で加算する値
    
    fn identity_map() -> Self::F {
        0 // 何もしない操作。乗算なら１とかね
    }

    fn mapping(&f: &Self::F, &x: &i64) -> i64 {
        x + f // 要素・区間最大値 x に操作 f を適用
    }

    fn composition(&f: &Self::F, &g: &Self::F) -> Self::F {
        f + g // 操作を合成
    }
}

let mut seg = LazySegtree::<RangeAddMax>::from(a);

seg.apply_range(1..4, 10); // a[1..4]の全要素に10を加算
seg.apply(2,5); // 一点に操作を適用
seg.set(2,100); // セグ木と同様、遅延せずに値そのものの更新も可能
seg.prod(0..3); // セグ木と同様に区間〇〇

&f: &Self::FはFの型を代入しているに近い。上の例なら&f: &i64と同義

FenwickTree

use ac_library::FenwickTree;

let mut fw = FenwickTree::new(n, 0_i64);

fw.add(2, 5);        // a[2] += 5
let sum = fw.sum(1..4); // a[1] + a[2] + a[3]

ModInt

use ac_library::ModInt998244353 as Mint;

let a = Mint::new(10);
let b = Mint::new(3);

let x = a / b;
println!("{}", x.val());

最大流

use ac_library::MfGraph;

let mut graph = MfGraph::<i64>::new(n);
graph.add_edge(from, to, capacity);

let ans = graph.flow(source, sink);

オーバーフローしそうなとき

// 配列範囲の左端を0に抑える
let left = i.saturating_sub(d);

// オーバーフローしたか判定したい
let result = a.checked_add(b);

// mod 2^N として意図的に回り込ませる
let result = a.wrapping_add(b);

座面圧縮

let mut xs = a.clone();
xs.sort_unstable();
xs.dedup();

let compressed: Vec<usize> = a
    .iter()
    .map(|&x| xs.partition_point(|&y| y < x))
    .collect();

今後追加するもの

SccGraph           // 強連結成分分解
TwoSat             // 2-SAT
MinCostFlowGraph   // 最小費用流

convolution        // 畳み込み
z_algorithm        // Z-algorithm
suffix_array       // 接尾辞配列
lcp_array          // LCP配列

pow_mod            // mod累乗
inv_mod            // mod逆元
crt                // 中国剰余定理
floor_sum          // floor和