・0 pv・3 min read
RustとAtCoderを勉強する(typical90_at)
はじめに
AtCoder の問題を Rust で解いていきます。AtCoder も Rust も初心者ですが、温かい目で成長を見守っていただけるとありがたいです。
今回は、競プロ典型90問の046 - I Love 46（★3）を解きました。
問題
つの長さ の数列 から、 を満たす が何通りあるかを求める問題です。
制約
- 入力は全て整数
提出コード
use itertools::iproduct;
use proconio::input;
use std::collections::HashMap;
fn vec2hashmap_mod46_counter(x: Vec<usize>) -> HashMap<usize, usize> {
x.into_iter()
.map(|x| x % 46)
.fold(HashMap::new(), |mut hashmap, item| {
let counter = hashmap.entry(item).or_insert(0);
*counter += 1;
hashmap
})
}
fn main() {
input! {
n: usize,
a: [usize; n],
b: [usize; n],
c: [usize; n],
}
let a = vec2hashmap_mod46_counter(a);
let b = vec2hashmap_mod46_counter(b);
let c = vec2hashmap_mod46_counter(c);
let mut ans = 0;
for (x, y, z) in iproduct!(a.keys(), b.keys(), c.keys()) {
if (x + y + z) % 46 == 0 {
ans += a[x] * b[y] * c[z];
}
}
println!("{}", ans);
}解説
となる性質を利用するそうです。まず、数列 の全要素に対して で割った余りを求めます。次に、それら余りが何個あるかを数えます。余りとその個数を求め、
HashMap にする関数を vec2hashmap_mod46_counter() としました。この関数を動作確認用に少し改変したものが下のコードです。
use std::collections::HashMap;
fn vec2hashmap_mod5_counter(x: Vec<usize>) -> HashMap<usize, usize> {
x.into_iter()
.map(|x| x % 5)
.fold(HashMap::new(), |mut hashmap, item| {
let counter = hashmap.entry(item).or_insert(0);
*counter += 1;
hashmap
})
}
let a = (1..9).collect();
println!("a: {:?}", &a);
let counter = vec2hashmap_mod5_counter(a);
println!("{:#?}", counter);a: [1, 2, 3, 4, 5, 6, 7, 8]
{
0: 1,
1: 2,
4: 1,
3: 2,
2: 2,
}
このように
HashMap<余り, 個数> とすれば、あとはlet mut ans = 0;
for (x, y, z) in iproduct!(a.keys(), b.keys(), c.keys()) {
if (x + y + z) % 46 == 0 {
ans += a[x] * b[y] * c[z];
}
}のようにして答えを求めることが出来ます。
まとめ
- 余り・合同式難しい
itertools::iproduct!が便利