【Rust】ABC429まとめ（A~D）

AtCoder Beginner Contest 429

後日バーチャル参加してA~Dの4完でした。
今回はA~Dをまとめていきます。

A - Too Many Requests

問題文の通り、$i \le M$のときOKを、そうでないとき、Too Many Requestsを出力してください。
さくっと書いて次に行きましょう。

use proconio::input;
use itertools::Itertools;

fn main() {
    input! {
        n: usize,
        m: usize,
    }

    let mut ans = vec![];
    for i in 0..n {
        if i < m {
            ans.push("OK")
        }else {
            ans.push("Too Many Requests")
        }
    }
    println!("{}", ans.iter().join("\n"));
}

B - N - 1

$\sum A$をあらかじめ求めておいて、$(\sum A) - A_i=M$となるような$i$が存在するかを調べます。
色々解法はあると思いますが、これが最もシンプルな気がします。

use proconio::input;

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

    let sum = a.iter().sum::<usize>();
    for &e in &a {
        if sum - e == m {
            return println!("Yes");
        }
    }
    println!("No");
}

C - Odd One Subsequence

同じ数字を2つ選び、残りの1つはその数字以外から選ぶ組み合わせを求める必要がありますが、これは数字ごとの出現回数がわかれば簡単に求めることができます。

長さ$N$の配列$A$の配列中に数字$x$が$M$回($M \ge 2)$回出現しているとすると、$x$の位置の組み合わせの数は$M(M-1)/2$で表せます。
そこに、$x$でない数字の個数$N-M$をかけた$M(M-1)(N-M)/2$が、$x$を2つ選んだ時の題意を満たす組み合わせの数となります。

この考え方で、$A$に2回以上登場する全ての数字について全探索して足し合わせたものが出力すべき答えになります。

use proconio::input;
use std::collections::HashMap;

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

    let mut freq: HashMap<usize, usize> = HashMap::new();
    for &e in &a {
        *freq.entry(e).or_insert(0) += 1;
    }

    let mut ans = 0;
    for (&_, &val) in &freq {
        if val < 2 {
            continue;
        } else {
            ans += (val * (val - 1) / 2) * (n - val);
        }
    }
    println!("{}", ans);
}

D - On AtCoder Conference

$i=0,\cdots, M-1$について、高橋くんを1マスずつ移動させて$X_i$を求めるようなやり方だと$O(M^2)$となり到底計算量が間に合いません。

ここで、工夫する点は大きく分けて二つあります。

まず、右端を調べる際に高橋くんを1マスずつ移動させるのではなく、人が存在する地点ごとに飛ばし飛ばしで移動させた方がよいです。人が存在しない地点をいくら動いても、区間内の人数は変わらないからです。
左端も同様にしたいのですが、求めたいのは$\sum^{M-1}_{i=0}$なので単に飛ばすことができません。
しかし、左端も人の存在する地点をまたがない限りは右端の位置や区間内の人数は変わらないので、$区間内の人数 \times (次に人の存在する地点 - 直前に人が存在した地点)$でまとめて処理できます。
以上で、全体計算量を$O(N^2)$まで落とすことができました。

次に、右端を調べる際に高橋くんを地点$i$からスタートさせるのではなく、前回の右端からスタートさせることができます。左端が単調増加のとき、区間内の人数を維持するには「その場にとどまる」か「右端を右にずらす」しかない（右端が単調非減少）からです。これは尺取法という考え方で、全体計算量を一気に$O(N)$まで落とせます。

▼尺取法による考え方のイメージ

上図ではさりげなく池の長さを二倍に伸ばして、人も同じ場所にコピーしていますが、このようにすることで右端の折り返しを考えなくて済むようになります。

以上の方針で、ソートがボトルネックとなり全体計算量$O(NlogN)$でこの問題を解くことができます。
実装については以下のACコードを参考にしてくださると幸いです。

use itertools::Itertools;
use proconio::input;
use std::collections::HashMap;

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

    let mut freq = HashMap::new();
    for &e in &a {
        *freq.entry(e).or_insert(0) += 1;
    }
    let sorted_set = freq.keys().copied().sorted().collect::<Vec<usize>>();
    let chained_a = sorted_set
        .iter()
        .chain(sorted_set.iter())
        .copied()
        .collect::<Vec<usize>>();
    let mut pointer = 1;
    let mut sum = 0;

    let mut ans = 0;
    for i in 1..=sorted_set.len() {
        while sum < c && pointer < 2 * n {
            sum += freq[&chained_a[pointer]];
            pointer += 1;
        }
        let (mut cur, prev) = (chained_a[i], chained_a[i - 1]);
        if chained_a[i] <= chained_a[i - 1] {
            cur += m;
        }
        ans += sum.min(n) * (cur - prev);
        sum -= freq[&chained_a[i]];
    }

    println!("{}", ans);
}

まとめ

CopilotなしでのRustのプログラミングにも慣れてきました。
また、次回コンテストから初期ユーザーの内部レートが800から1200に設定変更されるというアナウンスがありました。参考
ユーザーへの影響としては、以前より少しだけレートが上がりやすくなるそうです。
とはいえ、今参加するとどんどんレートが下がってしまいそうな感じがあります笑