ABC399:Rustで解く！問題解説

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

A問題

問題

https://atcoder.jp/contests/abc399/tasks/abc399_a

解説

文字列 $SS$ と $TT$ の対応する文字を1文字ずつ比較し、異なる文字の個数を数えます。

コード

abc399a.rs

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

fn main() {
    // 入力
    input! {
        n: usize,
        s: Chars,
        t: Chars
    }

    let mut cnt = 0;
    for i in 0..n {
        if s[i] != t[i] {
            cnt += 1;
        }
    }

    // 出力
    println!("{}", cnt);
}

B問題

問題

https://atcoder.jp/contests/abc399/tasks/abc399_b

解説

$NN$ 人がそれぞれ持っている得点 $PiP_i$ について、以下のルールで順位をつけます。

• 得点が大きい順に1位、2位、...とします。
• 同点の人が複数いる場合は、同点の人全員に同じ順位をつけます。

この問題を解くためには、以下1~3の手順を実行します。

1. 得点とその人のインデックスをペアにして、得点の降順にソートします。
2. ソートされた得点を順番に見ていき、前回の得点と比較します。
   • 前回の得点と異なる場合：現在のインデックス（1始まり）を順位として記録します。
   • 前回の得点と同じ場合：前回に記録した順位をそのまま適応します。
3. 最終的に、元の順序に対応する順位を出力します。

コード

abc399b.rs

use proconio::input;
const INF: usize = 1 << 60;

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

    // 値とインデックスをペアにして降順ソート
    let mut val_idx = Vec::new();
    for i in 0..n {
        val_idx.push((p[i], i));
    }
    val_idx.sort_by(|a, b| b.0.cmp(&a.0));

    let mut ans = vec![0; n];
    let mut prev_val = INF;
    let mut prev_rank = INF;

    for i in 0..n {
        if prev_val != val_idx[i].0 {
            // 前回値と異なる場合は、現在の順位を記録
            ans[val_idx[i].1] = i + 1;

            // 前回値と順位を更新
            prev_val = val_idx[i].0;
            prev_rank = i + 1;
        } else {
            // 前回値と一致する場合は、前回の順位を記録
            ans[val_idx[i].1] = prev_rank;
        }
    }

    // 出力
    for i in 0..n {
        println!("{}", ans[i]);
    }
}

C問題

問題

https://atcoder.jp/contests/abc399/tasks/abc399_c

解説

Union-Findというデータ構造を用いることで解くことができます。Union-Findは、頂点同士が同じ集合に属しているかを効率的に判定し、必要に応じて集合を統合することができるデータ構造です。

問題では、与えられた辺を1本ずつ確認し、以下1~2の操作を行います。

1. 辺が結ぶ頂点 $uu$ と $vv$ がすでに同じ集合の場合：閉路を防ぐためその辺を無視します。
2. 辺が結ぶ頂点 $uu$ と $vv$ が異なる集合の場合：その辺を用いて2つの集合を統合します。

無視した辺の本数が答えとなります。

コード

abc399c.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 ans = 0;
    for _ in 0..m {
        input! {
            u: Usize1, v: Usize1
        }
        // 頂点 u と v がすでに同じ集合に属している場合、無視する
        if !uf.merge(u, v) {
            ans += 1;
        }
    }

    // 出力
    println!("{}", ans);
}

// Union-Find
struct UnionFind {
    size: Vec<usize>,   // 各頂点が属する集合のサイズ
    parent: Vec<isize>, // 各頂点の親情報（根の場合は -1）
}

impl UnionFind {
    // 初期化
    fn new(n: usize) -> Self {
        UnionFind {
            size: vec![1; n],
            parent: vec![-1; n],
        }
    }

    // 頂点 u の代表元を求める（経路圧縮を行う）
    fn find(&mut self, u: usize) -> usize {
        if self.parent[u] == -1 {
            return u;
        }
        let p = self.parent[u] as usize;
        let root = self.find(p);
        self.parent[u] = root as isize; // 経路圧縮
        root
    }

    // 頂点 u と v を統合する（異なる集合の場合のみ統合）
    fn merge(&mut self, u: usize, v: usize) -> bool {
        let mut x = self.find(u);
        let mut y = self.find(v);

        // すでに同じ集合の場合は統合しない
        if x == y {
            return false;
        }

        // サイズが小さい集合を大きい集合に統合する
        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
    }

    // 頂点 u と v が同じ集合に属しているか判定
    fn same(&mut self, u: usize, v: usize) -> bool {
        self.find(u) == self.find(v)
    }

    // 頂点 u が属する集合のサイズを取得
    fn size(&mut self, u: usize) -> usize {
        let root = self.find(u);
        self.size[root]
    }
}

D問題

問題

https://atcoder.jp/contests/abc399/tasks/abc399_d

解説

問題文の以下が成り立つ条件を考えます。

$NN$ 組のカップルが一列に座っています。
$22$ 組のカップルであって、もともと両方のカップルは隣り合わせで座っておらず、かつ $44$ 人の間で席を交換することで両方のカップルが隣り合わせで座れるようになる組の個数を数えてください。

この内容を具体的にすると、以下1~3の条件を満たす組み合わせを探す問題となります。

1. 2組のカップルの値をそれぞれ $aa$ と $bb$ とします。
2. 各カップルの値が隣り合わせで座るように、席を交換する必要があります。
3. 入れ替え対象の値は、隣接する2つの値の間でのみ交換可能です。

例えば、以下のような並びが与えられたとします。

4 3 1 2 5 5 4 2 1 3

例えば、1つ目の 1 に注目すると、隣接する 2 や 3 と、もう1つの 1 と入れ替えることで、以下の状態を作り出すことができます。

• 「1のペア、2のペア」

4 3 1 1 5 5 4 2 2 3

• 「1のペア、3のペア」

4 1 1 2 5 5 4 2 3 3

一方で、「1のペア、4のペア」や「1のペア、5のペア」は、 4 や 5 は 1 と隣接していないため作り出すことができません。

したがって、隣接する2つの値に注目して交換することによって、条件を満たす組み合わせを数えます。実装としては以下1~4の通りになります。

1. 各値の出現位置を記録します。
2. 隣接する2つの値のペアを列挙します。ただし、値が同じ場合はスキップします。
3. 列挙したペアについて、交換によって条件を満たすかを確認します。
4. 条件を満たすペアの数をカウントします。

コード

abc399d.rs

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

fn main() {
    // テストケース数を取得
    input! {
        t: usize,
    }

    for _ in 0..t {
        solve();
    }
}

fn solve() {
    // 入力
    input! {
        n: usize,
        a: [Usize1; n * 2], // 各値のインデックスを0始まりで扱う
    }

    // 各値の出現位置を記録
    let mut position = vec![Vec::new(); n];
    for i in 0..2 * n {
        position[a[i]].push(i);
    }

    // 隣接する2つの値のペアをセットに記録
    let mut pair_set = HashSet::new();
    for i in 0..2 * n - 1 {
        // 隣接する2つの値が同じ場合はスキップ
        if a[i] == a[i + 1] {
            continue;
        }

        let mut x = a[i];
        let mut y = a[i + 1];
        if x > y {
            swap(&mut x, &mut y);
        }
        pair_set.insert((x, y));
    }

    let mut ans = 0;
    for &(x, y) in &pair_set {
        // xとyの出現位置を取得
        let x1 = position[x][0];
        let x2 = position[x][1];
        let y1 = position[y][0];
        let y2 = position[y][1];

        // 同じ値が隣接している場合はスキップ
        if x1 + 1 == x2 || y1 + 1 == y2 {
            continue;
        }

        // xとyを入れ替えた結果、xとyが隣接する場合をカウント
        if y1.abs_diff(x1) == 1 && y2.abs_diff(x2) == 1 {
            ans += 1;
        }
    }

    // 結果を出力
    println!("{}", ans);
}