ABC383:Rustで解く！問題解説

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

A問題

• 問題

https://atcoder.jp/contests/abc383/tasks/abc383_a

• 解説

現在の時刻 now_timeから次の時刻tまでの時間差分のリットル数の水を減少し、その後、水をvリットル追加します。この操作をn回繰り返し、最終的に残る水の量を出力します。

• コード

use proconio::input;
use std::cmp::max;

fn main() {
    input!{
        n: usize,
    }
    
    let mut now_time = 0; // 現在の時刻
    let mut now_l = 0; // 現在の水のリットル数
    
    for _ in 0..n {
        input!{
            t: i64, // 入力される時刻
            v: i64, // 入力される水のリットル数
        }
        // 現在時刻から次の時刻までの差分
        let diff = t - now_time;
        // 時刻の差分だけ水が減る（水量が0未満にならないようにする）
        now_l = max(0, now_l - diff);
        // 新たに水を補充する
        now_l += v;       
        // 現在の時刻を更新
        now_time = t;
    }
    // 最終的な水の量を出力
    println!("{}", now_l);
}

B問題

• 問題

https://atcoder.jp/contests/abc383/tasks/abc383_b

• 解説

.になっているマスを2つ選び、加湿器を仮置きして考えます。
1つ目の加湿器については、距離がD以内にあるマスをHで塗りつぶします。
2つ目の加湿器については、距離がD以内で1つ目の加湿器で塗りつぶしていないマスをHで塗りつぶします。2つの加湿器をチェックした際に塗りつぶしたマスの数をカウントします。
.のマス2箇所の選び方を全て試し、その結果の最大値を求めます。

• コード

use proconio::{input, marker::Chars};
use std::cmp::max;

fn main() {
    input! {
        h: usize, w: usize, d: usize,
        s: [Chars; h],
    }

    let mut ans = 0;

    // 置く場所を2つ選ぶ
    for i in 0..h * w {
        for j in 0..h * w {
            // 同じ場所2つでない、かつ'#'でない
            if i == j { continue; }
            let pos1 = (i / w, i % w);
            let pos2 = (j / w, j % w);
            if s[pos1.0][pos1.1] == '#' { continue; }
            if s[pos2.0][pos2.1] == '#' { continue; }

            // 実際に置いてみて温められるマスをカウント
            let mut cnt = 0; // 塗りつぶしたマスの数をカウント
            let mut ss = s.clone();
            
            // pos1, pos2から温め実施
            change_heat(&mut ss, pos1, &mut cnt, h, w, d);
            change_heat(&mut ss, pos2, &mut cnt, h, w, d);
            
            // 塗りつぶしたマスのmaxを記録
            ans = max(ans, cnt);
        }
    }
    println!("{}", ans);
}

fn change_heat(ss: &mut Vec<Vec<char>>, pos: (usize, usize), cnt: &mut usize, h: usize, w: usize, d: usize) {
    // 全てのマスをチェック
    for i in 0..h * w {
        // '#'でないマスを選択
        let chk_pos = (i / w, i % w);
        if ss[chk_pos.0][chk_pos.1] != '.' { continue; }
        
        // 距離がd以内かどうかを判定
        let dist = pos.0.abs_diff(chk_pos.0) + pos.1.abs_diff(chk_pos.1);       
        if dist > d { continue; }

        // 一度チェックしたマスを'H'にする
        ss[chk_pos.0][chk_pos.1] = 'H';
        *cnt += 1; // 塗りつぶしたマスをカウント
    }
}

C問題

• 問題

https://atcoder.jp/contests/abc383/tasks/abc383_c

• 解説

加湿器が置かれているマスHを全て開始地点として、BFS（幅優先探索）を行います。開始地点からの距離D以内で、上下左右に移動できるマスの個数を数えます。

• コード

use proconio::{input, marker::Chars};
use std::collections::VecDeque;
const INF: usize = 1 << 60;
const D4: [(usize, usize); 4] = [(!0, 0), (0, !0), (1, 0), (0, 1)]; // 上左下右

fn main() {
    input! {
        h: usize,    // 行数
        w: usize,    // 列数
        d: usize,    // 距離制限
        s: [Chars; h], // マスの状態を表す2次元配列
    }

    // 各マスの距離を初期化
    let mut dist = vec![vec![INF; w]; h];

    // Hのマスをすべてキューに積む
    let mut dq = VecDeque::new();
    for i in 0..h {
        for j in 0..w {
            if s[i][j] == 'H' {
                dist[i][j] = 0; // Hのマスの距離は0
                dq.push_back((i, j)); // キューに追加
            }
        }
    }

    // BFS
    while !dq.is_empty() {
        let cpos = dq.pop_front().unwrap(); // 現在の位置を取得
        // 4方向に移動
        for &dir in &D4 {
            let nx = cpos.0.wrapping_add(dir.0); // 新しい行
            let ny = cpos.1.wrapping_add(dir.1); // 新しい列
            // 範囲外チェックと壁チェック
            if nx >= h || ny >= w || s[nx][ny] == '#' { 
                continue; 
            }

            // 未訪問のマスの距離を更新
            if dist[nx][ny] == INF {
                dist[nx][ny] = dist[cpos.0][cpos.1] + 1;
                dq.push_back((nx, ny)); // 新しい位置をキューに追加
            }
        }
    }

    // distがd以下の個数を数える
    let mut ans = 0;
    for i in 0..h {
        for j in 0..w {
            if dist[i][j] <= d { 
                ans += 1; // 条件を満たすマスのカウント
            }
        }
    }
    println!("{}", ans); // 結果を出力
}

D問題

• 問題

https://atcoder.jp/contests/abc383/tasks/abc383_d

• 解説

まず、約数の個数がちょうど9個になる整数Xについて考えます。
これはXを素因数分解して、$X = P_1^{a_1} P_2^{a_2} ... P_m^{a_m}$という形に表した場合、
$(1 + a)(1 + b)...(1 + d) = 9$となる必要があります。($P_i$は素因数、$a_i$は任意の正の整数。)
ここで、上の式が成り立つ条件を考えると、以下の(1)(2)のどちらかを満たす時に成り立ちます。
(1)素因数が1つ、かつ$a_1$が8の場合
→ $X = P_1^8$ ($X = 2^8=256, 3^8=6561$ などが該当。)
(2)素因数が2つ、かつ$a_1, a_2$が2の場合
→ $X = P_1^2 P_2^2$ ($X = 2^2 \times 3^2 = 36, 2^2 \times 5^2 = 100$ などが該当。)

次に、(1)(2)の素因数を調べる範囲について考えます。
(1)の場合は、

$$\begin{aligned} P_1^8 & \le 4 \times 10^{12} \\ P_1 & \le \sqrt[8]{4 \times 10^{12}} \\ & \simeq 40 \end{aligned}$$

$P_1$は40以下となります。
また、(2)の場合は、

$$\begin{aligned} (P_1P_2)^2 & \le 4 \times 10^{12} \\ P_1P_2 & \le 2 \times 10^{6} \end{aligned}$$

$P_1,P_2$はそれぞれ$2 \times 10^6$以下です。
したがって、$2 \times 10^6$までの素数に対して、(1)(2)のいずれかを満たす値の個数を数え上げればよいです。

• コード

use proconio::input;
use libm::{pow, sqrt};

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

    // 約数がちょうど9個となるのは、以下の(1)(2)のどちらかを満たす時に成り立つ
    // (1)素因数が1個で、指数が8の場合
    // x^8 = 4*10^12 → xは40以下
    // (2)素因数が2個で、各々の指数が2の場合
    // x^2*y^2 = 4*10^12
    // xy = 2*10^6 → xまたはyの最大値は2*10^6

    // 素数列挙
    let sqrt_n = sqrt(n as f64) as usize;
    let erato = prime_list(sqrt_n);

    let mut ans = 0;

    // (1)の数え上げ
    for i in 0..erato.len() {
        let val = pow(erato[i] as f64, 8.) as usize;
        if val > n { break; }
        ans += 1;
    }

    // (2)の数え上げ
    for i in 0..erato.len() {
        for j in i + 1..erato.len() {
            let sqrt_val = erato[i] * erato[j];
            if sqrt_val > sqrt_n { break; }
            ans += 1;
        }
    }
    
    println!("{}", ans);
}

// エラトステネスの篩による素数リスト
pub fn prime_list(x: usize) -> Vec<usize> {
    let mut bools = vec![true; x + 1];
    bools[0] = false;
    bools[1] = false;
    let mut i = 2;
    while i * i <= x {
        if bools[i] {
            let mut now = 2 * i;
            while now <= x {
                bools[now] = false;
                now += i;
            }
        }
        i += 1;
    }
    let mut ret = Vec::new();
    for i in 0..=x {
        if bools[i] {
            ret.push(i);
        }
    }
    ret
}

E問題

• 問題

https://atcoder.jp/contests/abc383/tasks/abc383_e

• 解説

頂点xから頂点yに移動する際、なるべく辺の重みが小さいものを選んで通るため、辺の重みが小さい順に辺をつないで全域木を作ります。この全域木はクラスカル法のアルゴリズムを用いて構築出来ます。
2つの頂点を辺でつなぐ際に、数列Aにある頂点と数列Bにある頂点が初めて行き来できるようになるタイミングで、その頂点をペアにします。もし数列A同士や数列B同士で辺をつなぐ場合は、まだペアにしていない個数を合算し、ペアが作れるタイミングでまとめてペアにします。作成したペアの個数と辺の重みを掛け合わせて合計することで、
$\sum_{i=1}^{k} f(A_i, B_i)$の最小値を求めることができます。

• コード

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

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

    // 各頂点におけるマッチングの個数を数える
    let mut match_a = vec![0; n];
    let mut match_b = vec![0; n];
    
    // 入力に基づいてマッチング数を累積
    for i in 0..k {
        match_a[a[i]] += 1;
        match_b[b[i]] += 1;
    }

    // 辺の重みが小さい順に並べる
    let mut wuv = Vec::new();
    for (u, v, w) in uvw {
        wuv.push((w, u, v));
    }
    wuv.sort();

    let mut ans = 0;

    // 重みの小さい辺から順に全域木を作成
    // 連結時にAとBでペアにできるものがあれば、マッチングさせる。
    let mut uf = UnionFind::new(n, true, true);
    for &(w, u, v) in &wuv {
        // 既に同じグループの場合はスキップ
        if uf.same(u, v) { continue; }

        // uとvの親を特定
        let leader_u = uf.find(u);
        let leader_v = uf.find(v);

        // AのグループuとBのグループvでマッチング
        if match_a[leader_u] > 0 && match_b[leader_v] > 0 {
            let cnt = min(match_a[leader_u], match_b[leader_v]);
            ans += cnt * w;
            match_a[leader_u] -= cnt;
            match_b[leader_v] -= cnt;
        }
        // AのグループvとBのグループuでマッチング
        if match_a[leader_v] > 0 && match_b[leader_u] > 0 {
            let cnt = min(match_a[leader_v], match_b[leader_u]);
            ans += cnt * w;
            match_a[leader_v] -= cnt;
            match_b[leader_u] -= cnt;
        }

        // 辺をマージする
        uf.merge(u, v);

        // マージ後の親にマッチングの個数を合算
        match_a[uf.find(u)] = match_a[leader_u] + match_a[leader_v];
        match_b[uf.find(v)] = match_b[leader_u] + match_b[leader_v];
    }
    println!("{}", ans);
}

struct UnionFind {
    size: Vec<usize>,        // 各頂点のグラフサイズ
    parent: Vec<isize>,      // 親情報
    pass_cmp: bool,          // 経路圧縮の有無
    ubs: bool,               // Union-by-sizeの有無
}

impl UnionFind {
    // 経路圧縮とUnion-by-sizeを使用した初期化
    fn new(n: usize, pass_cmp: bool, ubs: bool) -> Self {
        UnionFind { size: vec![1; n], parent: vec![-1; n], pass_cmp, ubs }
    }

    // 代表値を求め、経路圧縮も行う
    fn find(&mut self, u: usize) -> usize {
        let mut cpos = u;
        if self.parent[cpos] == -1 { return cpos; }

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

        // 経路圧縮ありの場合
        if self.pass_cmp {
            // 親を根に付け替える
            let par = self.find(p);
            self.parent[cpos] = par as isize;
            return self.parent[cpos] as usize;
        } else {
            // 経路圧縮なしの場合、親を返す
            return self.find(p);
        }        
    }

    // 併合
    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.ubs {
            // 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
    }

    // 同じ集合かどうかを判定
    fn same(&mut self, u: usize, v: usize) -> bool {
        self.find(u) == self.find(v)
    }

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