Zenn
Open15

Rust で競技プログラミングをするときの細かいテクニック

Rust
#
競技プログラミング
#
競プロ
cunitaccunitac

while let を使った DFS や BFS

let mut stack = vec![start];
while let Some(a) = stack.pop() {
  for &b in &adjacent[a] {
    stack.push(b);
  }
}
cunitaccunitac

.iter().enumerate() が煩わしいときの itertools::izip

use itertools::izip;
for (i, &a) in izip!(0.., &a) {}
  • 1.. にしたりするのも簡単
  • for (i, (&a, &b)) in a.iter().zip(&b).enumerate() も for (i, &a, &b) in izip!(0.., &a, &b) と書ける
  • i が usize と推論されないこともあるので、0usize.. としておくと安心
cunitaccunitac

グリッド上の隣接点を得る

const N1: usize = 1usize.wrapping_neg();
const D4: [(usize, usize); 4] = [(1, 0), (0, 1), (N1, 0), (0, N1)];
fn adjacent4(
    (i, j): (usize, usize),
    (h, w): (usize, usize),
) -> impl Iterator<Item = (usize, usize)> {
    D4.iter()
        .map(move |&(di, dj)| (i.wrapping_add(di), j.wrapping_add(dj)))
        .filter(move |&(i, j)| i < h && j < w)
}
  • 必要なら impl Iterator<...> のあとに + でいろいろ書くとよい
  • 隣接8点なら [(1, 0), (1, 1), ...]
cunitaccunitac

グリッドのインデックスを簡単にする

use proconio::derive_readable;
use std::ops::{Index, IndexMut};

#[derive_readable]
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
struct Ix(usize, usize);

impl<T> Index<Ix> for Vec<Vec<T>> {
    type Output = T;
    fn index(&self, ix: Ix) -> &T {
        &self[ix.0][ix.1]
    }
}

impl<T> IndexMut<Ix> for Vec<Vec<T>> {
    fn index_mut(&mut self, ix: Ix) -> &mut T {
        &mut self[ix.0][ix.1]
    }
}
  • グリッドを使う長期コンテストとかでオススメ
  • 隣接点を得るやつを impl Ix に書いておくと便利
cunitaccunitac
for i in 0..h {
  // ここで何もしない
  for j in 0..w {
  // ここで何かする
  }
  // ここで何もしない
}

なら

use itertools::iproduct;
for (i, j) in iproduct!(0..h, 0..w) {
  // ここで何かする
}

でよい

cunitaccunitac

wrapping_neg

let a = 300usize;
let b = 45usize.wrapping_neg();
a.wrapping_add(b); // 255
cunitaccunitac

saturating_add など

let mut dist = vec![std::usize::MAX; 5];
dist[3] = 0;
for i in 1..5 {
  dist[i] = dist[i-1].saturating_add(10);
}
  • 謎の例になってしまったが、要するに初期値を MAX にしてもオーバーフローを気にせず計算できる
  • オーバーフローを気にしたほうが楽かも
  • usize だと思って saturating_sub を使っていたが、 実は i32 だったので普通に負になってて悲しい、みたいなこともある
cunitaccunitac

if let で「もし答えが存在しない場合は -1 を出力してください」をする

let ans: Option<Foo> = foo;
if let Some(ans) = ans {
  println!("{}", ans);
} else {
  println!("-1");
}
  • ans.unwrap_or(-1) とか ans.map_or(-1, |ans| ans as i32) できることもあるが、上の書き方のほうが気持ちいい……気がする
cunitaccunitac

数える HashMap

use std::collections::HashMap;

fn main() {
    let a = vec![1, 1, 2, 3, 3, 3];
    let b = vec![1, 1, 3];
    let mut count = HashMap::new();
    for &a in &a {
        // 追加
        *count.entry(a).or_insert(0usize) += 1;
    }
    for b in &b {
        // 削除
        let count_b = count.get_mut(b).unwrap();
        match *count_b {
            0 => panic!(),
            1 => { count.remove(b); }
            _ => *count_b -= 1,
        }
    }
}
  • 可変な参照を変数に束縛しておくのはわりと便利
  • よい感じのラッパを作ってもいいかも
  • もちろん BTreeMap でも
  • 別に remove 不要ならしなくてもよいし、0 => panic!() も不要なら *count.get_mut(b).unwrap() -= 1 で事足りる
cunitaccunitac

参照を持つ変数

let ai = &a[i];
let bi = &mut b[i];
  • ai や bi が Vec だったりすると特に便利
cunitaccunitac

0b11_111_111_111 は !(!0 << 11)

  • (1usize << 11).wrapping_neg(1) より気持ちいい
cunitaccunitac

インタラクティブな proconio

use proconio::{input, source::LineSource};
let stdin = std::io::stdin();
let mut stdin = LineSource::new(stdin.lock());
input!(from &mut stdin, a: i32);