Rust 競プロ チートシート集

Rustで競プロしているときに、これやりたいときどうやるだっけ??となりがちなので、チートシートを作りました。
適宜更新・修正します。
AtCoderにおける入出力方法は以下の記事を参照にしてください。

https://zenn.dev/tipstar0125/articles/f9c4626cdd4d5b

Vec

iter

- enumerate

forで配列をindexと一緒に回す。

let A = vec![3, 4, 5];
for (i, &a) in A.iter().enumerate() {
    println!("{} {}", i, a);
}
// 0 3
// 1 4
// 2 5

- rev

forで逆回し

let N = 3;
for i in (0..N).rev() {
    println!("{} ", i);
}
// 2 1 0
let A = vec![3, 4, 5];
for &a in A.iter().rev() {
    println!("{} ", a);
}
// 5 4 3

- enumerate + rev

let A = vec![3, 5, 2, 1, 4];
for (i, &a) in A.iter().enumerate().rev() {
    println!("{} {}", i, a);
}
// 4 4
// 3 1
// 2 2
// 1 5
// 0 3

let A = vec![3, 5, 2, 1, 4];
for (i, &a) in A.iter().rev().enumerate() {
    println!("{} {}", i, a);
}
// 0 4
// 1 1
// 2 2
// 3 5
// 4 3

- max/min

unwrap()忘れがち。空のときpanic。

let vec = vec![1, 2, 3];
println!("{}", vec.iter().max().unwrap());  // 3
println!("{}", vec.iter().min().unwrap());  // 1

- sum

型指定忘れがち。

let vec = vec![1, 2, 3];
println!("{}", vec.iter().sum::<usize>());  // 6

- swap

indexを指定して入れ替え。

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

変数同士のswapは以下の通り。

use std::mem::swap;  // 自動補完で入力されるように設定しておくと便利
let mut a = 0;
let mut b = 1;
swap(&mut a, &mut b);
println!("a:{}, b:{}", a, b);  // a: 1, b: 0

- extend

複数の配列を連結する。元の配列を再利用する場合は、cloneする。

let A = vec![1, 2, 3];
let B = vec![4, 5, 6];
let mut C = A;
C.extend(B);
println!("{:?}", C);  // [1, 2, 3, 4, 5, 6]

- join

配列を指定文字区切りで出力する。

A.iter().join(" ");  // スペース区切り
A.iter().join("\n");  // 改行区切り

- permutations

順列全探索

let n = 3;
for p in (0..n).permutations(n) {
    println!("{:?}", p);
}
// [0, 1, 2]
// [0, 2, 1]
// [1, 0, 2]
// [1, 2, 0]
// [2, 0, 1]
// [2, 1, 0]

// 以下でもOK
let vec = vec![0, 1, 2];
for p in vec.iter().permutations(vec.len()) {
    println!("{:?}", p);
}

next_permutationやprev_permutationもある。

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

- combinations

組合せの全パターンを列挙

for c in (0..3).combinations(2) {
    println!("{:?}", c);
}
// [0, 1]
// [0, 2]
// [1, 2]
let vec = vec![1, 2, 3, 3];
for c in vec.iter().combinations(2) {
    println!("{:?}", c);
}
// [1, 2]
// [1, 3]
// [1, 3]
// [2, 3]
// [2, 3]
// [3, 3]

以下のようにするとbit全探索と同じことができる。

let n = 3;
for i in 1..=n {
    for c in (0..n).combinations(i) {
        println!("{:?}", c);
    }
}
// [0]
// [1]
// [2]
// [0, 1]
// [0, 2]
// [1, 2]
// [0, 1, 2]

ソート

• 昇順、降順

let mut A = vec![3, 1, 2];
A.sort(); // 昇順
A.reverse();  // 降順
A.sort_by(|a, b| a.cmp(b)); // 昇順
A.sort_by(|a, b| b.cmp(a)); // 降順

• キーを指定してソート

let mut A = vec![(1, 3), (2, 1), (3, 2)];
A.sort_by(|(_, a), (_, b)| b.cmp(a)); // 第2要素で降順

let mut B = vec![("a", 90, 80), ("d", 70, 80), ("c", 90, 60), ("b", 70, 80)];
B.sort_by(|(a0, a1, a2), (b0, b1, b2)| (-a1, a2, a0).cmp(&(-b1, b2, b0)));
// 第2要素で降順、第3要素で昇順、第1要素で昇順の優先度で並び替え
// [("c", 90, 60), ("a", 90, 80), ("b", 70, 80), ("d", 70, 80)]

// マイナスを使用すると、型をisizeにする必要があるが、
// 以下のようにReverseを使用することでusizeのまま降順にできる。
use std::cmp::Reverse;
B.sort_by(|(a0, a1, a2), (b0, b1, b2)| (Reverse(a1), a2, a0).cmp(&(Reverse(b1), b2, b0)));

• 重複削除

Vecの重複を削除するときは事前にソートする。
ソートしないと、削除されない。。。

let mut A = vec![3, 1, 2];
A.sort();
A.dedup();

• 浮動小数点のソート

浮動小数点はNaNの可能性があるのでsort()メソッドは使用できない。NaNにならない（0割りがない）ことを自分で保証することを前提として、以下の通りソートすることをができる。

let mut A = vec![10.5, 2.3, 3.5];
//A.sort(); NG
A.sort_by(|a, b| a.partial_cmp(b).unwrap());
println!("{:?}", A);  // [2.3, 3.5, 10.5]

Vecの変換

use std::collections::VecDeque;
let mut A: VecDeque<_> = A.iter().cloned().collect();
let mut set: BTreeSet<_> = A.iter().cloned().collect();

スライス

例えば、答えを保存した配列を0からではなく、1からスペース区切りで出力したい場合配下の通り、Rangeで指定する。

let ans = vec![1, 2, 3, 4];  // 2 3 4
println!("{}", ans[1..].iter().join(" "));  // 2 3

文字列

chars

Stringを1文字ずつ処理するには、chars()を使用する。

let S = "abc".to_string();
// let S = "abc"; &strのままでもOK
for c in S.chars() {
    print!("{} ", c);  // a b c
}

proconioのCharsは、Rust標準のCharsとは異なり、Vec<char>である。
Vec<char>を生成する場合は以下の通りとする。

let s = "abc".chars();  // Chars
let ss = "abc".chars().collect_vec();  // Vec<char>

String or &str-> num

let s = "1";  // &str
let num: usize = s.parse().unwrap();
println!("{}", num);  // 1
let s = "2".to_string();  // String
let num: usize = s.parse().unwrap();
println!("{}", num);  // 2

Vec<char> -> String

let C = ['a', 'b', 'c'];
let S = C.iter().collect::<String>();  // "abc"

文字反転

let S = "abc".to_string();
let S_rev = S.chars().rev().collect::<String>();
println!("{}", S_rev);  // cba

スライス

let S = "abc".to_string();
println!("{}", &S[1..]);  // bc
println!("{}", &S[..2]);  // ab

文字をabc...の順に数値に変換

let s = 'a';
println!("{}", s as u8 - b'a' + 1);  // 1
let s = 'B';
println!("{}", s as u8 - b'A' + 1);  // 2, 2023/02/14 modified

u8をcharに戻す場合は、charでキャストすればよい。

let s = 'A';
let s_num = s as u8;
println!("{}", s_num as char);  // A

BTreeSet/BTreeMap

同じようなことができるHashSetとHashMapがあるが、デバッグするときに、昇順に並んでいた方が分かりやすい、二分探索ができる、というメリットがあるので、BTreeSetとBTreeMapを使用する。

BTreeMapの初期化

初期化の際、型を定義しなくても、類推してくれるが、Pythonでいうdefaultdictのような使い方をするために、型定義する。

let mut map: BTreeMap<usize, usize> = BTreeMap::new();
for &a in &A {
    *map.entry(a).or_default() += 1;  // デリファレンスが必要, 初期値：0
}

let mut map: BTreeMap<usize, Vec<usize>> = BTreeMap::new();
for (i, &a) in A.iter().enumerate() {
    map.entry(a).or_default().push(i + 1);  // 初期値：空配列
}

先頭・末尾の取得

let mut set = BTreeSet::new();
set.insert(10);
set.insert(20);
set.insert(30);
println!("{}", set.iter().next().unwrap());  // 10
println!("{}", set.iter().next_back().unwrap());  // 30

let mut map: BTreeMap<usize, usize> = BTreeMap::new();
*map.entry(10).or_default() = 30;
*map.entry(20).or_default() = 20;
*map.entry(30).or_default() = 10;
println!("{:?}", map.iter().next().unwrap());  // (10, 30)
println!("{:?}", map.iter().next_back().unwrap());  // (30, 10)

リセット

BTreeSetやBTreeMapを空にするにはclear()を使用する。

set.clear();
map.clear();

二分探索

HashSetではlower_boundが使用できない。代わりにBTreeSetのrangeを使用する。以下はBTreeSetの例だが、BTreeMapのキーでも同じことができる。

let mut set = BTreeSet::new();
let value = set.range(x..).next().unwrap(); // lower_bound相当
let value = set.range(..x).next_back().unwrap();  // xより小さい要素の中で最大の要素
let value = set.range(x..=y);  // x以上y以下

集合の計算

計算後の型が変わらず、使いまわせる以下を使用する。

let mut set = BTreeSet::new();
let mut set1 = BTreeSet::new();
set1.insert(1);
set1.insert(2);
set1.insert(3);
let mut set2 = BTreeSet::new();
set2.insert(1);
set2.insert(2);
set2.insert(4);
set = &set1 | &set2;  // 和
println!("{:?}", set);  // {1, 2, 3, 4}
set = &set1 - &set2;  // 差
println!("{:?}", set);  // {3}
set = &set1 & &set2;  // 積
println!("{:?}", set);  // {1, 2}
set = &set1 ^ &set2;  // 対称差
println!("{:?}", set);  // {3, 4}

以下は計算後の型がHashSetやBTreeSetではなく、Unionとかになるので、使い道なさそう??
計算結果を使いまわしたい場合に都合が悪い。into_iter().copied().collect()で変換できるけど、面倒。

a.union(&b)  // 和
a.difference(&b)  // 差
a.intersection(&b)  // 積
a.symmetric_difference(&b)  // 対称差

グラフ

グラフの頂点が文字列や大きな数値のときの変換

グラフの頂点が文字列や大きな数値のときは、隣接リストを配列ではなく、HashMapで作ることで処理できるが、以下のように変換するHashMapをつくり、隣接リストを配列で作ることができる。キーはユニークなので、同じキーのものは削除される。

let mut x = vec![];
for (s, t) in &ST {
    x.push(s.as_str());
    x.push(t.as_str());
}

let map: BTreeMap<_, _> = x.iter().enumerate().map(|(i, &s)| (s, i)).collect();

嵌りがちなエラー

• subtract with overflow

型がusizeのとき、引き算を使用した結果、負の値になるとエラーが発生する。isizeを使用するか、移項して引き算を使用しないように工夫する必要がある。

input! {
    N: usize,
    A: [usize; N]
}

for i in 0..N {
    // 負の値になるような計算をさせない
    if a - b > 0 {} // NG
    if a > b {} // OK
    
    // 負にならないことを確認してからindex指定する
    if A[a-b] > c {}  //NG
    if a > b && A[a-b] > c {}  //OK
    
    // 配列のindexを指定する変数の型がisizeの場合は、asでusizeに変更
    A[d as usize]
}

• 型推論で、型がi32とかになる場合、オーバーフローにより、WAになるので、型を指定

i32が扱える数の上限を超えてしまう場合、アルゴリズムは正しくでもWAになってしまい、何が原因でWAが出ているのか分からなくなるのを防ぐために、意識しておいた方が良い。また、なぜWAになったのか原因不明な場合は、変数にホバーして、意図していない型になっていないか確認する。

// これだけだと型推論でi32になる場合がある
let mut ans = 0; // NG(厳密にはこれで問題ない場合もある)
let mut ans: usize = 0; // OK
let mut ans = 0_usize; // OK

浮動小数点

三角関数

let theta = 180.0;
println!("{}", theta.to_radians());  // 3.141592653589793
let theta = theta.to_radians();
println!("{}", theta.to_degrees());  // 180

let theta = (180.0).to_radians();
println!("{}", theta.sin());  // 0.00000000000000012246467991473532
println!("{}", theta.cos());  // -1

その他

タプル分割代入

input! {
    N: usize,
    LR: [(usize, usize); N]
}

for i in 0..N {
    let l = LR[i].0;
    let r = LR[i].1;
    // 上述でも良いが、以下のように分割代入可能
    let (l, r) = LR[i];
    let (mut l, mut r) = LR[i]; // 値を変更する場合
}

パターンマッチ

// queryの第一要素でクエリのパターン(1 or 2 or 3)を指定し、その他の要素（引数）で処理
for &q in &query {
    match q {
        (1, x, y) => {
            // something
        }
        (2, _, _) => {
            // something
        }
        (3, x, _) => {
            // something
        }
        (_, _, _) => unreachable!(),
    }
}

大きい数の設定

1,000,000,007とかの0の数の大きい設定は以下の通り、省略して書ける。

let MOD = 1e9 as usize + 7;