yiskw note

機械学習やプログラミングについて気まぐれで書きます

【Rust】HashMapを用いて要素の個数を数える

rust

AtCoderのB - Judge Status SummaryをRustで解いている際に，
Rustでの配列内の要素の数え方がわからなかったので， 調べた内容をこちらにメモしておきます．
間違っていたり，別の方法がある場合はコメントにて教えていただけると幸いです．

HashMap

要素の数を数える方法には，Pythonでいう辞書型を用いて，要素とその個数を保持してあげれば良いです．
Rustではstd::collections::HashMapがそれに該当します．
ただPythonの辞書型と異なり，キーが存在しない場合の処理が面倒です．
Pythonではgetメソッドでキーが存在しない場合のデフォルト値を指定できますが，
RustではNoneしか返すことができません．
またPythonのdefaultdictのようなものもなさそうなので，キーが存在しない場合の処理は自分で書く必要があります．

具体的には以下のようにします．(参考)

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);

entryメソッドで，Entry型が戻り値として得られます．
これは列挙型(enum)で，キーに紐づいた値が存在するかどうかを表します．
or_insertメソッドによって，もしキーchに対応する値がある場合は，その値取り出されます．
キーchに対応する値がが存在しなければ，or_insertメソッドの引数の値がHashMapに追加されます．
この際or_insertの返り値は，可変参照になっていることに注意します．

問題の解法

以上を踏まえて，B - Judge Status Summaryは以下のように解きました．

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

fn main() {
    input! {
        n: usize,
        s: [String; n],
    };

    use std::collections::HashMap;
    let mut map = HashMap::new();

    for res in s {
        let counter = map.entry(res).or_insert(0);
        *counter += 1;
    }

    println!("AC x {}", map.get("AC").unwrap_or(&0));
    println!("WA x {}", map.get("WA").unwrap_or(&0));
    println!("TLE x {}", map.get("TLE").unwrap_or(&0));
    println!("RE x {}", map.get("RE").unwrap_or(&0));
}

参考