AtCoderでよく使いそうなitertoolsのメソッド

October 14, 2021

AtCoder では Rust 言語での提出で itertools が使えますが、今までほとんど使っていなかったので便利そうなメソッドをまとめました。AtCoder では itertools の0.9.0を利用できるので、このバージョンに含まれるメソッドのみを対象にしています。なお、AtCoder で使用できるクレートの一覧はこちらから見ることができます。

別の型への変換

イテレータ => ベクタ

let x = (1..=5).collect_vec(); // vec![1,2,3,4,5]

イテレータ => タプル

let tuple: Option<(usize, usize)> = (0..2).collect_tuple(); // Some((0,1))

イテレータ => イテレータ

dropping(n): 前の $n$ 個をスキップ

let iter = (1..3).dropping(1);
assert_eq!(iter, (2..3));

// これは標準のskipと同じ
let iter = (1..3).skip(1);
assert_eq!(iter, (2..3));

dropping_back(n): 後ろの $n$ 個をスキップ

let iter = (1..3).dropping_back(1);
assert_eq!(iter, (1..2));

combinations(r): $n$ 個から $r$ 個を選ぶ組み合わせ ( ${}_n C_r$ )

let iter = (1..=3).combinations(2);
assert_eq!(iter.collect_vec(), vec![[1, 2], [1, 3], [2, 3]]);

permutations(r): $n$ 個から $r$ 個を選ぶ順列 ( ${}_n P_r$ )

let iter = (1..=3).permutations(2);
assert_eq!(
    iter.collect_vec(),
    vec![[1, 2], [1, 3], [2, 1], [2, 3], [3, 1], [3, 2]]
);

cartesian_product(iter): 直積集合

let iter = (1..=2).cartesian_product(3..=4);
assert_eq!(iter.collect_vec(), vec![(1, 3), (1, 4), (2, 3), (2, 4)]);

// 2つのイテレータの型が異なっていても使える
let iter = (1..=2).cartesian_product(vec!['x', 'y']);
assert_eq!(
    iter.collect_vec(),
    vec![(1, 'x'), (1, 'y'), (2, 'x'), (2, 'y')]
);

interleave(iter): イテレータを交互に混ぜる

let iter = (1..=4).interleave(100..=101);
assert_eq!(iter.collect_vec(), vec![1, 100, 2, 101, 3, 4]);

// interleave_shortestの場合は交互にできなくなった時点で切り上げる
let iter = (1..=4).interleave_shortest(100..=101);
assert_eq!(iter.collect_vec(), vec![1, 100, 2, 101, 3]);

intersperse(x): $x$ を各要素の間に挿入

let iter = (1..=4).intersperse(0);
assert_eq!(iter.collect_vec(), vec![1, 0, 2, 0, 3, 0, 4]);

positions(f): ${\rm f}(x)$ を満たす $x$ のインデックスのみを取り出す

let iter = (1..=4).positions(|x| x % 2 == 0);
assert_eq!(iter.collect_vec(), vec![1, 3]); // 2と4のインデックス

unique(): 重複した値を取り除く

let iter = vec![1, 1, 2, 3, 2, 1].into_iter().unique();
assert_eq!(iter.collect_vec(), vec![1, 2, 3]);

// unique_byならf(x)が同一となるxを同一視する
let iter = vec![1, 1, 2, 3, 2, 1].into_iter().unique_by(|x| x % 2);
assert_eq!(iter.collect_vec(), vec![1, 2]);

dedup(): 隣接する重複した値を取り除く

// stdにはVecのdedup()はあるがイテレータにはない
let iter = vec![1, 1, 2, 3, 2, 1].into_iter().dedup();
assert_eq!(iter.collect_vec(), vec![1, 2, 3, 2, 1]);

// dedup_byならf(x,y)がtrueとなるxとyを同一視する
let iter = vec![1, 2, 4, 3, 1, 1]
  .into_iter()
  .dedup_by(|x: &i32, y: &i32| (x - y).abs() < 2);
assert_eq!(iter.collect_vec(), vec![1, 4, 1]);

その他便利メソッド

join(): イテレータから文字列生成

競プロでは配列を標準出力するのに使える。

let formatted = (0..3).join(" ");
println!("{}", formatted); // 0 1 2

concat(): ベクタの合成

let vectors = vec![vec![1], vec![2, 3], vec![4, 5, 6]];
assert_eq!(vectors.into_iter().concat(), vec![1, 2, 3, 4, 5, 6]);