ABC392:Rustで解く！問題解説

AtCoder Beginner Contest 392のA~F問題をRustで解いた際の解法をまとめました。

A問題

問題

https://atcoder.jp/contests/ABC392/tasks/ABC392_a

解説

整数列 $A$ を昇順に並び替えたとき、条件 $A_1 \times A_2 = A_3$ が成り立つかを判定します。

コード

abc392a.rs

use proconio::input;

fn main() {
    input!{
        mut a: [usize; 3],
    }   
    a.sort();
    
    println!("{}", if a[0] * a[1] == a[2] {"Yes"} else {"No"});
}

B問題

問題

https://atcoder.jp/contests/ABC392/tasks/ABC392_b

解説

$1$ から $N$ の要素集合から、整数列 $A$ に含まれる要素集合を取り除いた差集合を求め、以下の2つの情報を出力します。

• 差集合の要素数
• 差集合の各要素

コード

abc392b.rs

use proconio::input;
use std::{collections::BTreeSet, iter::FromIterator};
use itertools::Itertools;

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

    // 1 ~ Nの集合1
    let mut st_1 = BTreeSet::new();
    for i in 1..=n { st_1.insert(i); }

    // 数列aの集合2
    let st_2 = BTreeSet::from_iter(a.iter().cloned());

    // 差集合
    let st = &st_1 - &st_2;

    // 差集合の要素数と各要素を出力
    println!("{}", st.len());
    println!("{}", st.iter().join(" "));
}

C問題

問題

https://atcoder.jp/contests/ABC392/tasks/ABC392_c

解説

問題文で求められている以下の内容について考えます。

i が書かれたゼッケンを着けている人が見つめている人の着けているゼッケンにかかれている数

上記を整理すると、以下の3点がポイントになります。これらを順番に処理していきます。

1. ゼッケン $i$ を着けている人 $x$
2. 人 $x$ が見つめている人 $y$
3. 人 $y$ が着けているゼッケンの番号 $z$

これに対し、人 $i$ について入力から読み取れる情報は次のようになります。

• 人 $i$ が見つめている人 $P_i$ → 上記の2番目に相当
• 人 $i$ が着けているゼッケンの番号 $Q_i$ → 上記の3番目に相当

ここで、入力の内容を逆に見ると、さらに以下のことが言えます。

• 人 $P_i$ を見ている人 $i$
• ゼッケンの番号 $Q_i$ を着けている人 $i$ → 上記の1番目に相当

したがって、人 $i$ ごとの1~3の変換結果を事前に記録した上で、順番に適用することが可能です。

コード

abc392c.rs

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

fn main() {
    input! {
        n: usize, 
        p: [Usize1; n],
        q: [Usize1; n],
    }

    // ①ゼッケンの番号Q_iを着けている人i
    let mut zekken_info = vec![0; n];
    for i in 0..n { 
        zekken_info[q[i]] = i; 
    }
    
    // ②人iが見つめている人P_i
    let mut stare_info = vec![0; n];
    for i in 0..n { 
        stare_info[i] = p[i]; 
    }

    // ③人iが着けているゼッケンの番号Q_i
    let mut wear_info = vec![0; n];
    for i in 0..n { 
        wear_info[i] = q[i]; 
    }

    // ①→②→③と変換
    let mut ans = Vec::new();
    for i in 0..n {
        let mut val = zekken_info[i];
        val = stare_info[val];
        val = wear_info[val];
        ans.push(val + 1);
    }

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

D問題

問題

https://atcoder.jp/contests/ABC392/tasks/ABC392_d

解説

各サイコロについて、$1$ から $10^5$ の目が出る確率を事前に計算しておきます。その後、サイコロ2個を選んで、ぞろ目の確率の合計を求めます。
これをすべてのパターンで試し、その中で最も大きい値を出力します。

コード

abc392d.rs

use proconio::input;
const SZ: usize = 100_000;

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

    // 確率リスト
    // px_list[サイコロi個目][出目j] = 確率
    let mut px_list = vec![vec![0.0; SZ + 1]; n];
    for i in 0..n {
        input! {
            k: usize,
            a: [usize; k]
        }
        px_list[i] = get_px_data(k, &a);
    }

    // サイコロ2個を全探索
    // ぞろ目確率の合計を計算し、最大値を求める
    let mut ans = 0.0;
    for i in 0..n {
        for j in i + 1..n {
            let tot_px = calc_px(i, j, &px_list);
            ans = f64_max(ans, tot_px);
        }
    }
    println!("{}", ans);
}

fn get_px_data(k: usize, a: &Vec<usize>) -> Vec<f64> {
    // 出目ごとの個数を数える
    let mut cnt = vec![0; SZ + 1];
    for &aa in a {
        cnt[aa] += 1;
    }

    // 出目ごとの確率に変換する
    let mut px_data = vec![0.0; SZ + 1];
    for i in 1..=SZ {
        px_data[i] = cnt[i] as f64 / k as f64;
    }
    px_data
}

fn calc_px(i: usize, j: usize, px_list: &Vec<Vec<f64>>) -> f64 {
    // ぞろ目の確率を合計
    let mut tot = 0.0;
    for k in 0..=SZ {
        tot += px_list[i][k] * px_list[j][k];
    }
    tot
}

fn f64_max(a: f64, b: f64) -> f64 {
    if a >= b { return a; }
    b
}

E問題

問題

https://atcoder.jp/contests/ABC392/tasks/ABC392_e

解説

サーバーを頂点、ケーブルを辺とするグラフを考え、すべての頂点を行き来できる全域木を作成します。全域木の作成にはUnion-Findというデータ構造を使用します。

まず、 $M$ 本の辺の2点間について、連結する2つの頂点が異なるグループであればマージを行い、同じグループであれば後でマージするための辺として記録します。その後、記録しておいた辺を1本ずつ参照し、辺の始点から見て①「異なるグループの頂点」が見つかるまで探して、追加でマージします。
なお、①については、頂点 $1$ から頂点 $N$ までの全体を1度だけ確認すれば十分です。

コード

abc392e.rs

use proconio::{input, marker::Usize1};
use std::mem::swap;

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

    // Union-Find 構造体の初期化
    let mut uf = UnionFind::new(n);

    // 連結不要な辺を記録する(辺の番号、始点)
    let mut edge = Vec::new();

    // 連結が必要な辺をすべてマージする
    for i in 0..m {
        input!{
            u: Usize1, v: Usize1,
        }
        // 同じグループであれば不要な辺を記録する
        if uf.same(u, v) {
            edge.push((i, u));
        }
        // 異なるグループであればマージする
        else {
            uf.merge(u, v);
        }
    }

    // 既に見た頂点を管理するための変数
    let mut epos = 0;
    // 追加でマージした辺の情報を保存するベクター
    let mut merge_edge_info = Vec::new();

    // 記録した辺を使用してマージを行う
    for (idx, spos) in edge {
        // 全頂点が全域木になるか、マージが必要な頂点が見つかるまでループ
        while epos < n && uf.same(spos, epos) { epos += 1; }

        // 全域木になっている場合はループを抜ける
        if epos == n { break; }

        // マージを行う
        uf.merge(spos, epos);
        merge_edge_info.push((idx, spos, epos));
    }

    // 追加でマージした辺の本数と、その辺の情報を出力する
    println!("{}", merge_edge_info.len());
    for (idx, spos, epos) in merge_edge_info {
        println!("{} {} {}", idx + 1, spos + 1, epos + 1);
    }
}

struct UnionFind {
    size: Vec<usize>,        // 各頂点のグラフサイズ
    parent: Vec<isize>,      // 親情報
}

impl UnionFind {
    fn new(n: usize) -> Self {
        UnionFind { size: vec![1; n], parent: vec![-1; n] }
    }
    
    // 代表値を求める。(経路圧縮も行う)
    #[allow(unused)]
    fn find(&mut self, u: usize) -> usize {
        let mut mut_u = u;
        if self.parent[mut_u] == -1 { return mut_u; }

        let p = self.parent[mut_u] as usize;

        // xの親を根に付け替える
        let par = self.find(p);
        self.parent[mut_u] = par as isize;
        return self.parent[mut_u] as usize;
    }

    // 併合
    fn merge(&mut self, u: usize, v: usize) -> bool {
        // 代表値を取得
        let mut x = self.find(u);
        let mut y = self.find(v);
        
        // 既に同じ集合であれば false を返す
        if x == y { return false; }
        
        // 木のサイズが小さいものを大きいものに併合する (Union by size)
        if self.size[x] < self.size[y] {
            swap(&mut x, &mut y);
        }
        
        // yをxの子とする
        self.parent[y] = x as isize;
        self.size[x] += self.size[y];
        true
    }

    // 2つの頂点が同じ集合かどうかを判定
    #[allow(unused)]
    fn same(&mut self, u: usize, v: usize) -> bool {
        self.find(u) == self.find(v)
    }

    // 集合のサイズを求める
    #[allow(unused)]
    fn size(&mut self, u: usize) -> usize {
        let x = self.find(u);
        self.size[x]
    }
}

F問題

問題

https://atcoder.jp/contests/ABC392/tasks/ABC392_f

解説

配列に要素を insert で一つずつ追加するだけの単純な問題に見えますが、insert() 1回あたりの計算量は $O(N)$ であり、$N$ 回実行すると合計で $O(N^2)$ になります。
このため、実行時間制限に間に合わないため、他の方法で解く必要があります。

具体的には、挿入操作の順番を逆にし、「値が未格納の配列」の先頭から $P_i$ 番目に値を入れる動作を繰り返す方法を採用します。ここで、 $P_i$ の個数は配列の先頭から見ると単調増加になっているため、二分探索の考え方が使えます。また、$P_i$ の個数の数え方はセグメントツリーで高速化し、先頭から $X$ 個見たときに、その区間内に $P_i$ 番目が存在するかどうかを特定し、挿入位置として適切な位置を見つけます。

二分探索とセグメントツリーの計算量はそれぞれ $O(\log{N})$ であり、 $N$ 回実行すると合計で $O(N \log{N})$ となるため、実行時間制限内で解くことが可能になります。

コード

abc392f.rs

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

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

    // セグメントツリーの初期化(0:値格納済み、1:値未格納)
    let mut sg = SegmentTree::new(n + 1, 0, combine_func);
    for i in 1..=n { sg.set(i, 1); }

    let mut ans = vec![0; n + 1];

    // pを後ろから処理する。
    for i in (1..=n).rev() {
        // セグメントツリー上を二分探索して、
        // 未格納の数が丁度p[i]個の位置を見つける
        let pos = find_lower_bound(&mut sg, 1, n, p[i - 1]);
        ans[pos] = i;

        // p[i]を格納済みにする。
        sg.set(pos, 0);
    }

    // 配列の2番目以降を出力
    println!("{}", ans.iter().skip(1).join(" "));
}

// セグメントツリーライブラリ
struct SegmentTree<F, T> {
    length: usize,  // 入力の長さ
    n: usize,       // 列の長さ
    array: Vec<T>,  // 二分木
    identity_e: T,  // 単位元(0:合計、max、1:積、INF:min)
    combine_f: F,   // 評価関数
}

impl<F: Fn(T, T) -> T, T: Copy + std::fmt::Display> SegmentTree<F, T> {
    // 初期化
    #[allow(unused)]
    fn new(length: usize, identity_e: T, combine_f: F) -> Self {
        let mut n = 1;
        while n < length { n <<= 1; }
        SegmentTree { length, n, array: vec![identity_e; 2 * n], identity_e, combine_f }
    }

    // [1点更新] 位置 index (0-indexed) を値 value で更新
    #[allow(unused)]
    fn set(&mut self, index: usize, val: T) {
        let mut i = self.n + index;
        self.array[i] = val;
        while i > 1 {
            i >>= 1;
            self.array[i] = (self.combine_f)(
                self.array[i << 1 | 0],
                self.array[i << 1 | 1],
            );
        }
    }

    // [1点取得] 位置 index (0-indexed) 内の値を取得
    #[allow(unused)]
    fn get(&mut self, index: usize) -> T {
        let mut l = index + self.n;
        let mut r = index + 1 + self.n;
        let mut val_l = self.identity_e;
        let mut val_r = self.identity_e;
        while l < r {
            if l & 1 != 0 {
                val_l = (self.combine_f)(val_l, self.array[l]);
                l += 1;
            }
            if r & 1 != 0 {
                r -= 1;
                val_r = (self.combine_f)(self.array[r], val_r);
            }
            l >>= 1; r >>= 1;
        }
        (self.combine_f)(val_l, val_r)
    }

    // [区間取得] 区間 [l, r) (0-indexed) 内の要素について
    // l 番目から順に、 combine_f を適応した結果を返す
    #[allow(unused)]
    fn prod(&mut self, lindex: usize, rindex: usize) -> T {
        let mut l = lindex + self.n;
        let mut r = rindex + self.n;
        let mut val_l = self.identity_e;
        let mut val_r = self.identity_e;
        while l < r {
            if l & 1 != 0 {
                val_l = (self.combine_f)(val_l, self.array[l]);
                l += 1;
            }
            if r & 1 != 0 {
                r -= 1;
                val_r = (self.combine_f)(self.array[r], val_r);
            }
            l >>= 1; r >>= 1;
        }
        (self.combine_f)(val_l, val_r)
    }

    #[allow(unused)]
    pub fn dbg_print(&mut self) {
        for i in 0..self.length {
            print!("{} ", self.get(i));
        }
        println!();
    }
}

// 評価関数 
pub fn combine_func(a: usize, b: usize) -> usize { 
    a + b
}

#[allow(unused)]
// セグメントツリー上の二分探索
fn find_lower_bound(
    sg: &mut SegmentTree<impl Fn(usize, usize) -> usize, usize>, // セグメントツリー
    lpos: usize,    // 左閉区間
    rpos: usize,    // 右開区間
    val: usize      // 値
) -> usize {
    let mut ng = lpos - 1;
    let mut ok = rpos;
    while ok - ng > 1 {
        let mid = (ok + ng) / 2;
        // 右開区間で取得する
        let x = sg.prod(lpos, mid + 1);
        if x >= val { 
            ok = mid; 
        } else { 
            ng = mid; 
        }
    }
    ok
}