5か月前公開・5か月前更新・3 min read

RustとAtCoderを勉強する(typical90_at)

https://cdn.apollon.ai/media/notebox/63b22052-8bc0-4145-95b3-a26daae60f7d.jpeg

はじめに

AtCoder の問題を Rust で解いていきます。AtCoder も Rust も初心者ですが、温かい目で成長を見守っていただけるとありがたいです。

問題

3

つの長さ

N

の数列

A=(A_1,A_2,\cdots,A_N),B=(B_1,B_2,\cdots,B_N),C=(C_1,C_2,\cdots,C_N)

から、

A_i+B_j+C_k\equiv0(mod\,46)

を満たす

(i,j,k)(1 \le i,j,k \le N)

が何通りあるかを求める問題です。

制約

$1 \le N \le 10^5$
$0 \le A_i,B_j,C_k \le 10^9$
入力は全て整数

提出コード

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

解説

(A_i+B_j+C_k)(mod\,46)=((A_i \, mod \, 46)+(B_j \, mod \, 46)+(C_k \, mod \, 46))(mod\,46)

となる性質を利用するそうです。まず、数列

A,B,C

の全要素に対して

46

で割った余りを求めます。次に、それら余りが何個あるかを数えます。余りとその個数を求め、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];
    }
}

RustとAtCoderを勉強する(typical90_at)

はじめに

問題

制約

提出コード

解説

まとめ

Discussion

コメントにはログインが必要です。