Magicode logo
Magicode
0
3 min read

RustとAtCoderを勉強する(typical90_at)

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

はじめに

AtCoder の問題を Rust で解いていきます。AtCoder も Rust も初心者ですが、温かい目で成長を見守っていただけるとありがたいです。
今回は、競プロ典型90問046 - I Love 46(★3)を解きました。

問題

33 つの長さ NN の数列 A=(A1,A2,,AN),B=(B1,B2,,BN),C=(C1,C2,,CN)A=(A_1,A_2,\cdots,A_N),B=(B_1,B_2,\cdots,B_N),C=(C_1,C_2,\cdots,C_N) から、Ai+Bj+Ck0(mod46)A_i+B_j+C_k\equiv0(mod\,46) を満たす (i,j,k)(1i,j,kN)(i,j,k)(1 \le i,j,k \le N) が何通りあるかを求める問題です。

制約

  • 1N1051 \le N \le 10^5
  • 0Ai,Bj,Ck1090 \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);
}

解説

(Ai+Bj+Ck)(mod46)=((Aimod46)+(Bjmod46)+(Ckmod46))(mod46)(A_i+B_j+C_k)(mod\,46)=((A_i \, mod \, 46)+(B_j \, mod \, 46)+(C_k \, mod \, 46))(mod\,46) となる性質を利用するそうです。まず、数列 A,B,CA,B,C の全要素に対して 4646 で割った余りを求めます。次に、それら余りが何個あるかを数えます。余りとその個数を求め、HashMap にする関数を vec2hashmap_mod46_counter() としました。この関数を動作確認用に少し改変したものが下のコードです。
rust
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! が便利

Discussion

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