Rustで配列内の要素の個数をカウントする

July 17, 2020

atcoderでCODE FESTIVAL 2017 qual BのB問題Problem Setを解いている時に、配列の要素の個数をカウントしたかったのですが、方法がわからなかったので調べたものをまとめます。

配列の要素の個数を数えたい時はHashMapを使用する

Rustの標準ライブラリにあるハッシュ集合(HashMap)を使用することで簡単に配列の要素の個数をカウントすることができます。

実装方法はドキュメントにも書かれている通りです。j

use std::collections::HashMap;

let mut letters = HashMap::new();

for ch in "a short treatise on fungi".chars() {
    let counter = letters.entry(ch).or_insert(0);
    *counter += 1;
}

assert_eq!(letters[&'s'], 2);
assert_eq!(letters[&'t'], 3);
assert_eq!(letters[&'u'], 1);
assert_eq!(letters.get(&'y'), None);

この部分がポイント

    let counter = letters.entry(ch).or_insert(0);
    *counter += 1;

この部分で、counterにはキーとなるchに対応する値がmutableな参照として渡されています。詳しくはor_insertの実装でご確認ください。 chという値の可変な参照として保存されているcounterの参照外しを行い、1を足しているのが個数のカウントの処理に値します。 or_inserはchというキーがなかった場合には0という値を挿入し、同様にその値をmutableな参照として渡しているのも注意です。

自分の答案, 環境

Rust(1.42.0) 標準入力にはproconioを使用しています。

#[allow(unused_imports)]
use proconio::marker::{Bytes, Chars, Usize1};
use proconio::{fastout, input};
use std::collections::HashMap;
#[fastout]
fn main() {
    input! {
        n: usize,
        d: [usize; n],
        m: usize,
        t: [usize; m],
    }
    let mut map_d: HashMap<usize, usize> = HashMap::new();
    for i in 0..n {
        let counter = map_d.entry(d[i]).or_insert(0);
        *counter += 1;
    }
    for i in 0..m {
        let counter = map_d.entry(t[i]).or_insert(0);
        if *counter > 0 {
            *counter -= 1;
        } else {
            println!("NO");
            return;
        }
    }
    println!("YES");
}

参考文献

Struct std::collections::HashMap Rust By Example 日本語版ハッシュ集合