緑コーダーがRustで解説してみた(ABC432 A～E)

AtCoder Beginner Contest 432のA～E問題を緑コーダーが自分なりの解説と解答コードをまとめました。
参考になりましたら幸いです。

ABC432-A

問題

https://atcoder.jp/contests/abc432/tasks/abc432_a

3桁の値を並び替えて、最も大きな値を作る問題です。

解説

与えられた3つの値を降順に並び替えることで、最も大きな値を作ることができます。
具体的には、3つの値を配列に格納し、降順にソートした後、それらを結合して出力します。

コード

abc432a.rs

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

fn main() {
    input! {
        mut abc: [usize; 3], // 3つの値
    }
    // 降順にソート
    abc.sort_by(|i, j| j.cmp(i));
    // ソートした値を結合して出力
    println!("{}", abc.iter().join(""));
}

ABC432-B

問題

https://atcoder.jp/contests/abc432/tasks/abc432_b

値の各桁を並び替えて、最も小さな値を作る問題です。

解説

与えられた値の各桁を並び替えた中で、最も小さな値を求めます。
ただし、先頭の桁が0であってはならないため、並び替えた結果の中で先頭が0でないものを選びます。

この条件を満たすためには、順列全探索を行い、先頭の桁が0でないものを探します。
条件を満たす最初に見つけた並びが最も小さな値となります。

コード

abc432b.rs

use itertools::Itertools;
use proconio::input;
use superslice::Ext;

fn main() {
    input! {
        x: usize, // 値(最大5桁)
    }
    // 値をchar配列に変換
    let mut cx: Vec<char> = x.to_string().chars().collect();
    // 最も小さい値を見つける
    find_minimize(&mut cx);
}

fn find_minimize(cx: &mut Vec<char>) {
    // 順列全探索して、先頭が0でないものを見つける
    cx.sort();
    while {
        if cx[0] != '0' {
            // 条件を満たす最小値を出力
            println!("{}", cx.iter().join(""));
            return;
        }
        // 次の順列を生成
        cx.next_permutation()
    } {}
}

ABC432-C

問題

https://atcoder.jp/contests/abc432/tasks/abc432_c

小さいアメと大きいアメを重さの総量が同じになるよう $N$ 人に配るときの、大きいアメの合計個数の最大値を求める問題です。

解説

大きいアメの合計個数を最大化するためには、以下1.～4.の通りに配るのが最適です。

1. 最初の人に全て大きいアメを配る
   最初の人には全て大きいアメを配り、その人のアメの重さの総量を基準とします。この基準に従って、他の人にもアメを配ります。

2. 他の人のアメの配り方を決定する
   他の人には、最初の人のアメの重さの総量と一致するように、大きいアメと小さいアメを組み合わせて配ります。

   • まず、全て大きいアメを配ると仮定し、その場合の重さを計算します。
   • その重さと基準との差を求めます。
   • 差を埋めるために、大きいアメを小さいアメに交換します。この交換が可能かどうかを判定し、可能ならば大きいアメの個数を計算します。

3. 交換条件の判定

   • 差が、大きいアメと小さいアメの重さの差の倍数である必要があります。
   • また、交換後の大きいアメの個数が負にならないようにします。

4. 結果の計算
   各人について、大きいアメの個数を計算し、その合計を求めます。途中で交換条件を満たさなせなかった場合は、-1 を出力します。

コード

abc432c.rs

use proconio::input;

fn main() {
    input! {
        n: usize, // 人数
        x: i64, // 小さいアメの重さ
        y: i64, // 大きいアメの重さ
        mut a: [i64; n], // 各人に配るアメの個数
    }
    // 個数の小さい順にソート
    a.sort();

    // 大きいアメの個数の合計
    let mut ans = 0;
    
	// 1人目には全て大きなアメにする
    ans += a[0];

    // 各人に配るアメの重さの総量を決定
    let weight = a[0] * y;

    // 2番目以降を決めていく
    for idx in 1..n {
        // 各人に配る大きなアメの個数を求める
        let cnt = decide_distribute(a[idx], x, y, weight);
        if cnt == -1 {
            println!("-1");
            return;
        }
        ans += cnt;
    }
    // 合計の大きいアメの個数を出力
    println!("{}", ans);
}

fn decide_distribute(tot_cnt: i64, x: i64, y: i64, weight: i64 ) -> i64 {
    // 全て大きなアメと仮定して、大きいアメ->小さいアメに何個変更するかを決める
    let assume_weight = tot_cnt * y;
    let diff_weight = assume_weight - weight;
    // 1個当たりの重さの差分
    let diff_unit = y - x;

    // 変更した際に重さが一致しなくなる場合
    if diff_weight % diff_unit != 0 {
        return -1;
    }

    // 小さなアメに交換する個数
    let small_cnt = diff_weight / diff_unit;
    let large_cnt = tot_cnt - small_cnt;
    
    // 大きなアメが0未満になる
    if large_cnt < 0 {
        return -1;
    }

    // 交換可能
    large_cnt
}

ABC432-D

問題

https://atcoder.jp/contests/abc432/tasks/abc432_d

ブロックが左右または上下にずれる操作を繰り返した後、各連結成分のブロックのサイズを答える問題です。

解説

ブロックの操作を実際にシミュレーションして解きます。
具体的には、以下1.～4.の手順で解くことができます。

1. シミュレーション
   $N$ 個のイベントを順に処理し、ブロックの操作をシミュレーションします。
   各イベントでは、指定された操作軸に基づいてブロックを移動させます。
   この時、ブロックが操作軸を横断する場合はブロックを2つに分割します。

2. 連結成分の構築
   シミュレーション後のブロックについて、重なる領域がある場合は連結します。
   この連結処理には、Union-Findというデータ構造を使用します。

3. 面積の計算
   各連結成分について、ブロックの面積を計算します。

4. 結果の出力
   連結成分の個数と、各連結成分のブロックの面積を昇順に並べて出力します。

コード

abc432d.rs

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

#[allow(unused)]
#[derive(PartialEq, Eq, PartialOrd, Ord, Hash, Clone)]
struct Rectange {
    xl: i64, xr: i64,
    yl: i64, yr: i64,
}

fn main() {
    input! {
        n: usize, // イベントの数
        x: i64, y: i64, // 初めのブロックサイズ
    }

    // 初めのブロックをセット
    let mut cur_block = Vec::new();
    cur_block.push(Rectange{ 
        xl: 0, xr: x, 
        yl: 0, yr: y
    });

    // イベントの実施
    for _ in 0..n {
        let mut next_block = Vec::new();
        input! {
            c: char, // 操作する方向(XまたはY)
            a: i64, // 操作軸の座標
            b: i64, // 左右または上下にずれる幅
        }

        // 各ブロックについて操作を行う
        for &Rectange { xl, xr, yl, yr } in &cur_block {
            if c == 'X' {
                // 操作軸の座標が間にある場合は、ブロックを2つに分裂
                if xl < a && a < xr {
                    add_separate_block(c, xl, a, xr, yl, yr, b, &mut next_block);
                }
                // 操作軸の座標より左側にある場合
                else if xr <= a {
                    add_shift_block(c, xl, xr, yl, yr, -b, &mut next_block);
                }
                // 操作軸の座標より右側にある場合
                else {
                    add_shift_block(c, xl, xr, yl, yr, b, &mut next_block);
                }
            }
            else {
                // 操作軸の座標が間にある場合は、ブロックを2つに分裂
                if yl < a && a < yr {
                    add_separate_block(c, xl, a, xr, yl, yr, b, &mut next_block);
                }
                // 操作軸の座標より上側にある場合
                else if yr <= a {
                    add_shift_block(c, xl, xr, yl, yr, -b, &mut next_block);
                }
                // 操作軸の座標より下側にある場合
                else {
                    add_shift_block(c, xl, xr, yl, yr, b, &mut next_block);
                }
            }
        }
        cur_block = next_block;
    }

    // 暫定のブロックの個数
    let cnt = cur_block.len();

    // ブロックの連結
    let mut uf = UnionFind::new(cnt);
    merge_area(&cur_block, cnt, &mut uf);

    // ブロック毎の面積を求める
    let mut area = get_component_area(&cur_block, cnt, &mut uf);

    // 面積が0のものを取り除く
    distinct_area_zero(&mut area);

    // 昇順に戻して、ブロックの個数と各面積を出力
    print_areas(&mut area);
}

fn add_separate_block(
    c: char, xl: i64, a: i64, xr: i64, yl: i64, yr: i64, 
    shift: i64, next_block: &mut Vec<Rectange>) 
{
    if c == 'X' {
        next_block.push(Rectange{
            xl: xl        , xr: a         ,
            yl: yl - shift, yr: yr - shift, 
        });
        next_block.push(Rectange{
            xl: a         , xr: xr        ,
            yl: yl + shift, yr: yr + shift, 
        });
    }
    else {
        next_block.push(Rectange{
            xl: xl - shift, xr: xr - shift, 
            yl: yl        , yr: a,
        });
        next_block.push(Rectange{
            xl: xl + shift, xr: xr + shift, 
            yl: a         , yr: yr        ,
        });
    }
}

fn add_shift_block(
    c: char, xl: i64, xr: i64, yl: i64, yr: i64, 
    shift: i64, next_block: &mut Vec<Rectange>) 
{
    if c == 'X' {
        next_block.push(Rectange{
            xl , xr,
            yl: yl + shift, yr: yr + shift, 
        });
    }
    else {
        next_block.push(Rectange{
            xl: xl + shift, xr: xr + shift, 
            yl , yr
        });
    }
}

fn merge_area(cur_block: &Vec<Rectange>, cnt: usize, uf: &mut UnionFind){
    // ブロック2つの組み合わせを調べる
    for a in 0..cnt {
        for b in a+1..cnt {
            let block_a = &cur_block[a];
            let block_b = &cur_block[b];
            
            // ブロックの隣接・重なりを調べる
            let check_xr = min(block_a.xr, block_b.xr);
            let check_xl = max(block_a.xl, block_b.xl);
            let check_yr = min(block_a.yr, block_b.yr);
            let check_yl = max(block_a.yl, block_b.yl);
            let cx = check_xr - check_xl;
            let cy = check_yr - check_yl;

            // 隣接・重なりがある場合はマージする
            if cx < 0 || cy < 0 { continue; }
            if cx > 0 || cy > 0 { uf.merge(a, b); }
        }
    }
}

fn get_component_area(cur_block: &Vec<Rectange>, cnt: usize, uf: &mut UnionFind) -> Vec<i64>{
    let mut area = vec![0; cnt];
    for a in 0..cnt {
        let block = &cur_block[a];
        let ret = (block.xr - block.xl) * (block.yr - block.yl);
        area[uf.find(a)] += ret;
    }
    area
}

fn distinct_area_zero(area: &mut Vec<i64>) {
    area.sort_by(|i, j| j.cmp(i));
    while *area.last().unwrap() == 0 {
        area.pop();
    }
}

fn print_areas(area: &mut Vec<i64>) {
    area.sort();
    println!("{}", area.len());
    for &mut val in area {
        print!("{} ", val);
    }
    println!();
}

// Union-Find構造体
struct UnionFind {
    size: Vec<usize>,        // 各頂点の集合サイズ
    parent: Vec<isize>,      // 親情報
}

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

    // 代表元を求める
    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
    }

    // 2つの集合を併合
    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);
        }
        self.parent[y] = x as isize;
        self.size[x] += self.size[y];
        true
    }

    // 同じ集合かどうかを判定
    #[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 root = self.find(u);
        self.size[root]
    }
}

ABC432-E

問題

https://atcoder.jp/contests/abc432/tasks/abc432_e

数列 $A$ の値を更新しながら、総和を求める問題です。

解説

この問題は、セグメントツリーというデータ構造を用いることで解くことができます。
セグメントツリーを使うことで、値の更新と区間の総和を効率的に行うことができます。

この問題では、以下2つのセグメントツリーを用意します。

• 個数セグメントツリー
  • 各値の出現回数を管理します。
• 合計値セグメントツリー
  • 各値の「値 × 出現回数」を管理します。

上のセグメントツリーを用意したうえで、次の2種類のクエリについて順番に処理します。

1. 値の更新

   • 数列 $A$ の特定のインデックスの値を変更します。
   • 変更前の値と変更後の値に基づいて、個数セグメントツリーと合計値セグメントツリーを更新します。

2. 区間の総和の計算

   • 入力で与えられる $l$ と $r$ に基づいて、各値を読み替えて総和を計算します。

   • 求める値は以下の図の通りになります。

     

     • $l > r$ の場合
       • 全ての値を $l$ とみなして、総和を計算します。
     • $l \leq r$ の場合、以下の3つの区間に分けて計算します。
       • 区間 $[0, l+1)$ : 全ての値を $l$ とみなして総和を計算します。
       • 区間 $[l+1, r+1)$ : 値そのものを合計します。
       • 区間 $[r+1, \infty)$ : 全ての値を $r$ とみなして総和を計算します。

コード

abc432e.rs

use proconio::{input, marker::Usize1};
const SZ: usize = 550_000;

fn main() {
    input! {
        n: usize, // 数列の長さ
        q: usize, // クエリ数
        mut a: [usize; n], // 数列の各値
    }

    // 数列について、値ごとの個数を求める
    let cnt = calc_cnt(&a);

    // 区間内の個数を取得するためのセグメントツリー
    let mut sg_cnt = SegmentTree::new(SZ, 0, combine_func);
    init_sg_cnt(&mut sg_cnt, &cnt);

    // 区間内の合計値を取得するためのセグメントツリー
    let mut sg_val = SegmentTree::new(SZ, 0, combine_func);
    init_sg_val(&mut sg_val, &cnt);

    // クエリ処理
    for _ in 0..q {
        input! {
            nm: usize, // クエリタイプ
        }
        if nm == 1 {
            input! {
                xi: Usize1, // 変更する数列のインデックス(0-index)
                y: usize, // 変更する値
            }
            let from = a[xi];
            let to = y;
            
            // 数列の値を更新
            a[xi] = y;

            // 個数の更新
            update_cnt(&mut sg_cnt, from, -1);
            update_cnt(&mut sg_cnt, to, 1);

            // 合計値の更新
            update_val(&mut sg_val, from, -(from as i64));
            update_val(&mut sg_val, to, to as i64);
        } else {
            input! {
                l: usize, // maxをとる値
                r: usize // minをとる値
            }
            if l > r {
                println!("{}", acc_type1(&mut sg_cnt, l, r));
            } else {
                println!("{}", acc_type2(&mut sg_cnt, &mut sg_val, l, r));
            }
        }
    }
}

fn calc_cnt(a: &Vec<usize>) -> Vec<i64> {
    let mut ret = vec![0; SZ];
    for &aa in a {
        ret[aa] += 1;
    }
    ret
}

fn init_sg_cnt<F>(
    sg: &mut SegmentTree<F, i64>, cnt: &Vec<i64>)
        where F: Fn(i64, i64) -> i64
{
    // 各値について、個数をセット
    for idx in 0..SZ {
        sg.set(idx, cnt[idx]);
    }
}
fn init_sg_val<F>(
    sg: &mut SegmentTree<F, i64>, cnt: &Vec<i64>)
        where F: Fn(i64, i64) -> i64
{
    // 各値について、値×個数をセット
    for idx in 0..SZ {
        sg.set(idx, idx as i64 * cnt[idx]);
    }
}

fn update_cnt<F>(
    sg: &mut SegmentTree<F, i64>, idx: usize, add: i64)
        where F: Fn(i64, i64) -> i64
{
    let from = sg.get(idx);
    let to = from + add;
    sg.set(idx, to);
}
fn update_val<F>(
    sg: &mut SegmentTree<F, i64>, idx: usize, add: i64)
        where F: Fn(i64, i64) -> i64
{
    let from = sg.get(idx);
    let to = from + add;
    sg.set(idx, to);
}

fn acc_type1<F>(
    sg_cnt: &mut SegmentTree<F, i64>, l: usize, _r: usize) -> i64
        where F: Fn(i64, i64) -> i64
{
    let mut ret = 0;
    // 全てmaxをとる区間
    ret += sg_cnt.prod(0, SZ) * l as i64;
    ret
}
fn acc_type2<F>(
    sg_cnt: &mut SegmentTree<F, i64>, 
    sg_val: &mut SegmentTree<F, i64>,     
    l: usize, r: usize) -> i64
        where F: Fn(i64, i64) -> i64
{
    let mut ret = 0;
    // maxをとる区間
    ret += sg_cnt.prod(0, l+1) * l as i64;
    // 総和の区間
    ret += sg_val.prod(l+1, r+1);
    // minをとる区間
    ret += sg_cnt.prod(r+1, SZ) * r as i64;
    ret
}


// セグメントツリーライブラリ
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)
    }
}
// 評価関数 
pub fn combine_func(a: i64, b: i64) -> i64 { 
    a + b
}