AtCoder Beginner Contest 352をRustで楽しむ

はじめに

AtCoder Beginner Contest 352をRustで解いてみました。

今回のコンテストはE問題まで解けたのでEまでの問題で私が解答したコードとコンテスト中に考えていたことをまとめます。

また、問題は各見出しのリンクからご覧ください。

私達ハウインターナショナルでは、社名をもじってハウッカソンという名のイベントを毎月最終金曜日に実施しています。

ハウッカソンのテーマに競技プログラミングを選ぶメンバーもいます。

A - AtCoder Line

私の解答

読解力がなさすぎて一回では理解できなかったですが、
よくよく読んでみるとXとYの間にZがありますか？というだけの問題だったことがわかります。

use proconio::input;
use std::cmp::{min, max};

fn main() {
    input! {
        _n: usize,
        x: usize,
        y: usize,
        z: usize
    }

    if min(x, y) <= z && max(x, y) >= z {
        println!("Yes");
    } else {
        println!("No");
    }
}

B - Typing

私の解答

Tを前からループしながらSの何文字目までが入力できているかをインクリメントして持っておけばよさそうです。
また、一致した時にansというvectorにその時のindexを1-indexedで格納して、それを出力しています。

use itertools::Itertools;
use proconio::{input, marker::Chars};

fn main() {
    input! {
        s: Chars,
        t: Chars
    }
    
    let mut now = 0usize;
    let mut ans = vec![];

    for (idx, &c) in t.iter().enumerate() {
        if c == s[now] {
            now += 1;
            ans.push(idx + 1);
        }
        if now == s.len() {
            break;
        }
    }

    println!("{}", ans.iter().join(" "));
}

C - Standing On The Shoulders

私の解答

愚直にすべての巨人の置き方を試そうとすると、O(N!)くらいになるので到底間に合いません。
ですが、ここで少し考えてみると、次のことがわかると思います。

• 頭の長さが気になるのは一番上に立っている巨人だけ
• 一番上以外の巨人たちの並び順はどう変えても高さは同じ

ということで、一番上にくる巨人を全探索すればよいです。
しかし、各探索で和をとっていてはO(N ^ 2)になってしまい、まだ実行時間の制限に間に合わないです。
そこで、あらかじめすべての巨人の肩までの高さの和 sum をとっておくことで、
sum - a + bを計算すれば、各巨人が一番上にいる時の高さがわかります。
この計算自体はO(1)でできるので、合計でO(N)で最大値を求められるというわけです。

ちなみに、ansはusizeでとって　ans = max(ans, sum - a + b) などで計算することでより高速に求められますが、自分はデバッグする際に println!("{:?}", ans) などですばやく途中計算を見たいのでベクタに入れてあとでmaxを取るように実装しがちです。

use proconio::input;

fn main() {
    input! {
        n: usize,
        ab: [(usize, usize); n]
    }
    
    let sum = ab.iter().map(|&e| e.0).sum::<usize>();
    let mut ans = vec![];
    for &(a, b) in ab.iter() {
        ans.push(sum - a + b);
    }
    println!("{}", ans.iter().max().unwrap());
}

D - Permutation Subsequence

私の解答

読解力がないため２回読みました。
そして、説明力もないので、入力1を用いて考えたことを説明してみます。

4 2
2 3 1 4

良い添字列というものになるにはPiは連続した整数列でなければならないので、
連続するようにソートしてあげます。
また、このとき、その数が配列Pの何番目にあるかを0-indexedで持たせておきます。
つまり、上記入力では以下のようになります。

[(1, 2), (2, 0), (3, 1), (4, 3)]

このような配列の隣合うk個の配列の中のインデックス番号の最大値と最小値の差の中でその差を最も小さくできる時の数を求めてねー。という問題です。
上記の例ではk=2なので、となりの要素との計算を行えばよいです。

例えば、0番目と1番目の要素を見てみると、

max(2, 0) = 2
min(2, 0) = 0

なので、この差は2です。

同じようにして、1番目と2番目の要素では 1 , 2番目と3番目の要素では 2 と計算でき、
これらの中で最小の 1 を出力してくださいという問題です。

上記の例ではk = 2で非常に小さいためなんだか愚直にやってもO(N)でとけそうな気がしてしまいますが、
制約を見てみると k <= 2 * 10 ^ 5 で非常に大きいため工夫が必要です。

ここで下の例で考えてみましょう。配列は適当、k = 3の例です。　
今、黒線で引かれた範囲の計算ができたとしましょう。

次に計算する範囲で変わる範囲としては

• 先頭を一つ追加(緑の要素を追加)
• 末尾を一つ削除(赤の要素を削除)

の高々2つしか変化がないことがわかります。
よって、この2回操作を高速に行うことで各ループごとに範囲を1から考えることなしに、計算したい範囲を決定できるということがわかりました。

あとは、この範囲の最大・最小値を高速に求められればよいですが、これは平衡二分探索木というデータ構造を用いれば実現できます。

Rustでは標準ライブラリに BTreeSet が用意されているので、これを使えば良いです。
setなので、ダブった値があったら飲み込まれてしまいますが、今回は、index番号なので、必ず一意な番号になることが保証されているので、BTreeSetでいけそうです。
よって、各範囲のインデックス番号最大・最小の取得もO(log|set|)と、高速に求められ、この問題が解けました！

use std::collections::BTreeSet;
use proconio::input;

fn main() {
    input! {
        n: usize,
        k: usize,
        p: [usize; n]
    }
    
    let mut v = vec![];
    for i in 0..n {
        v.push((p[i], i));
    }
    v.sort_by_key(|e| e.0);

    let mut set = BTreeSet::new();
    for i in 0..k - 1 {
        set.insert(v[i].1);
    }

    let mut ans = vec![];

    for i in 0..n - k + 1 {
        set.insert(v[i + k - 1].1);
        let min = set.iter().next().unwrap();
        let max = set.iter().last().unwrap();
        ans.push(max - min);

        set.remove(&v[i].1);
    }
    println!("{}", ans.iter().min().unwrap());
}

E - E - Clique Connect

私の解答

前置き・前提知識

最小全域木って問題に書いてくれていますね。優しいです。
おそらく、ここまで読んでくださっている方がそれなりにアルゴリズム知識がある方だと思っておりますので、最小全域木とは何かの説明は省きます。

最小全域木問題の解法には以下の3つがあります。

• クラスカル法
• プリム法
• ブルーフカ法

そして、今回自分はこの中のクラスカル法を使ってこの問題を解きました。

クラスカル法は

• 重みが小さい辺から貪欲に辺を追加していく(追加した辺の重みを足しておく)
• ただし、グラフにサイクルができる場合はスキップ

というめちゃくちゃ簡単なステップで最小全域木問題を解くことができます。
グラフにサイクルができるかどうかの判定はUnionFindを使えば良いですね。

問題の考察

さて、ここからがこの問題の考察ですが、部分集合Sのすべての頂点に対して辺を追加する処理があります。
つまり、 S = {1, 2, 3, 4, 5}だったら、以下のように辺を張らなければならないのです。

これを実際に張っていこうと思うと、各集合で kC2個の組み合わせが考えられ、到底間に合いません。
しかし、次のことを考えてみます。

• 部分集合の中の重みは全て均一
• グラフが連結になった瞬間に、それ以降はもう他の辺は最小全域木ステップの中で追加されない

ということで、考える辺というのはその部分集合のグラフを連結にできるような最小の辺の成分だけでよいことになり、これは、k - 1本ですね。
kの制約を見てみると、

$$
\sum k \le 2 \times 10^{5}
$$

なので、これであれば追加できます。
あとは上記手順に従って解くことでこの問題が解けました！

忘れがちですが、連結にならなかった時の判定は、処理がすべて終わった後に uf.groups()のlengthをみればよいです。
下記実装では、集合の0番目にそのほかの頂点を連結させています。

use ac_library::Dsu;
use proconio::{input, marker::Usize1};

fn main() {
    input! {
        n: usize,
        m: usize,
    }
    
    let mut v = vec![];
    for _ in 0..m {
        input! {
            k: usize,
            c: usize,
            a: [Usize1; k]
        }
        for i in 1..k {
            v.push((a[0], a[i], c));
        }
    }
    v.sort_by_key(|e| e.2);

    let mut uf = Dsu::new(n);
    let mut ans = 0usize;

    for &(n1, n2, cost) in v.iter() {
        if uf.same(n1, n2)  {
            continue;
        }
        uf.merge(n1, n2);
        ans += cost;
    }
    if uf.groups().len() != 1 {
        println!("-1");
    } else {
        println!("{}", ans);
    }
}

おわりに

コンテスト中にA ~ Eまでを解けたので、考えていたことを簡単に説明してみました。
Eまでのはや解きでなんとかレート上がりましたが黄Diff解けるようになりたいので精進します。