競プロ典型90問をRustで解いていこう（mod）13日目

2022-05-14

Rust

はじめに

本日の問題

競プロ典型90問をRustで解いていこう（準備）0日目

今日は上の記事から13問目を解いていきます。

046 - I Love 46（★3）を解いていこうと思います。

046 - I Love 46（★3）

同じ長さの数列が3つ与えられるので、それぞれの数列の値を足して46の倍数になる選び方の総数を求めるというものです。

Rustでの回答例

与えられる値の全探索だとタイムアウトするだろうかと思いながら試したものを置いときます。

use proconio::input;

fn main() {
    input! {
        n: usize,
        a: [u128;n],
        b: [u128;n],
        c: [u128;n]
    }

    let mut ans = 0;

    for i in 0..n {
        for j in 0..n {
            for k in 0..n {
                let e = (a[i] + b[j] + c[k]) % 46;
                if e == 0 {
                    ans += 1;
                }
            }
        }
    }
    println!("{}", ans);
}

上記方針では上手く行きませんでした。別の方法を考える必要がありそうです。

Rustでの回答例(OK)

総当りでやるのではなくある程度制約から対象を絞るようにやるのがいいそうです。今回であれば、0~46の余りを持つものが何個存在するのかを求めます。あとは、各iの個数の組み合わせ合計で得たい物が得られるはずと考えてこのようにしました。

use proconio::input;

fn main() {
    input! {
        n: usize,
        a: [usize;n],
        b: [usize;n],
        c: [usize;n]
    }

    let mut amod46 = vec![0; 46];
    let mut bmod46 = vec![0; 46];
    let mut cmod46 = vec![0; 46];
    for i in 0..n {
        amod46[a[i] % 46] += 1;
        bmod46[b[i] % 46] += 1;
        cmod46[c[i] % 46] += 1;
    }

    let mut ans: u128 = 0;
    for i in 0..46 {
        for j in 0..46 {
            for k in 0..46 {
                if (i + j + k) % 46 == 0 {
                    ans += amod46[i] * bmod46[j] * cmod46[k];
                }
            }
        }
    }
    println!("{}", ans);
}

これで通るはずです。