Pythonと比べて学ぶ競プロのためのRust Part2 #Python

今までは Python で競技プログラミングをしていましたが、最近 Rust の勉強を始めました。同じように Python を書いたことがあり Rust へ入門しようと考えている人向けに AtCoder の過去問 (B~C) を両言語で解いてみます。

この記事は Part2 です。

前回の記事

Pythonと比べて学ぶ競プロのためのRust

前回のテーマ

基本文法 ( for / if / 入出力)

参照

メモ化再帰

Option 型

HashMap / defaultdict / 辞書

今回扱うテーマ

累積和 / scan
オーバーフロー
all / any
lower_bound
座標圧縮

ABC238 B - Pizza

問題：

Python

提出：https://atcoder.jp/contests/abc238/submissions/31425241

I = lambda:int(input())
LI = lambda:list(map(int,input().split()))
from itertools import accumulate

n = I()
al = LI()

deg = [0] + sorted([a%360 for a in accumulate(al)]) + [360]
print(max(deg[i+1] - deg[i] for i in range(n+1)))

Rust

提出：https://atcoder.jp/contests/abc238/submissions/31421717

use itertools::{sorted, Itertools};

fn main() {
    proconio::input! {
        n:usize,al:[usize;n]
    }
    let deg = sorted(
        vec![
            vec![0, 360],
            al.iter().scan(0_usize, |s, x| {
                    *s = (*s + x) % 360;
                    Some(*s)
                }).collect_vec()]
        .concat()).collect_vec();
    println!("{}",
        deg.iter().tuple_windows().map(|(&a, &b)| b - a).max().unwrap()
    );
}

scan 関数を用いて al の累積和を作成する。 itertools の collect_vec によりベクタに変換する。
Part1 でも扱ったが、ベクタの連結は vec![vec1,vec2].concat() .
collect_vec() は collect::<Vec<_>>() と同じ。
scan は配列から新たな配列を作るときに用いる。
- 第一引数に初期値,第二引数に Some(_) 値を返す関数をとる。
- s の型は &mut なので値を用いるときは *s の形で使う。
- scan に似た関数に fold がある。これは配列から一つの値を求めるときに使う。
itertools の sorted を用いると変数を mut で宣言する必要がなくなる。
隣り合う切込みの差の最大値を求めるため tuple_windows を用いる。
- itertools のtuple_windows は隣り合う要素をタプルで取り出す関数。

ABC233 C - Product

問題：

Python

提出：https://atcoder.jp/contests/abc233/submissions/31439179

I = lambda:int(input())
LI = lambda:list(map(int,input().split()))
from collections import defaultdict

n,x = LI()
vl = defaultdict(int)
vl[1] = 1
for _ in [0]*n:
    l,*al = LI()
    _vl = defaultdict(int)
    for k,v in vl.items():
        for a in al:
            _vl[k*a] += v
    vl = _vl
print(vl[x])

Rust

提出：https://atcoder.jp/contests/abc233/submissions/31380953

use std::collections::HashMap;

fn main() {
    proconio::input! {
        n:usize,x:usize
    }
    let mut vl = HashMap::new();
    vl.insert(1usize, 1usize);
    for _ in 0..n {
        proconio::input! {
            l:usize,al:[usize;l]
        }
        vl = {
            let mut _vl = HashMap::new();
            for (k, v) in vl.iter() {
                for a in al.iter() {
                    *_vl.entry((*k).saturating_mul(*a)).or_insert(0) += *v;
                }
            }
            _vl
        };
    }
    println!("{}", *vl.entry(x).or_insert(0));
}

vl は key：ボールの総積, val : 取り出し方の総数の辞書。
HashMap を用いてPythonの defaultdict のように扱う方法は Part1 に書いたものと同様。
Python と違い Rust ではオーバーフローに気をつける必要がある。
- usize 型の最大値は $1.8\times 10^{19}$ 程度。
- a.saturating_mul(b) は a*b がオーバーフローしないなら a*b を、オーバーフローするなら usize の最大値を返す関数。
- 今回は使っていないが、同じくオーバーフロー対策で使える関数に a.overflowing_mul(b) がある。これは計算結果とオーバーフローしているかを表す bool のペアを返す関数。

ABC232 C - Graph Isomorphism

問題：

Python

提出：https://atcoder.jp/contests/abc232/submissions/31439077

I = lambda:int(input())
LI = lambda:list(map(int,input().split()))
from itertools import permutations

n,m = LI()
ab = [LI() for _ in [0]*m]
cd = [LI() for _ in [0]*m]
edge = set([])
for a,b in ab:
    edge.add((a-1,b-1))
    edge.add((b-1,a-1))
check = False
for p in permutations(range(n)):
    check |= all([(p[cd[i][0]-1],p[cd[i][1]-1]) in edge for i in range(m)])
print(["No","Yes"][check])

Rust

提出：https://atcoder.jp/contests/abc232/submissions/31381551

use std::collections::HashSet;

use itertools::Itertools;
use proconio::marker::Usize1;

fn main() {
    proconio::input! {
        n:usize,m:usize,
        ab:[(Usize1,Usize1);m],
        cd:[(Usize1,Usize1);m],
    }
    let mut edge = HashSet::new();
    for (a, b) in ab {
        edge.insert((a, b));
        edge.insert((b, a));
    }
    let mut check = false;
    for p in (0..n).permutations(n) {
        check |= (0..m).all(|i| edge.contains(&(p[cd[i].0], p[cd[i].1])));
    }
    println!("{}", ["No", "Yes"][check as usize]);
}

全ての点の並べ替えに対して繋ぎ方が同じであるかを判定する。
入力時に Usize1 型で受け取ると 0-index での usize 型の値として受け取ることができる。便利。
HashSet は HashMap の値に固定値 () を指定して作られている集合型。和集合や差集合を求めるなどの集合に対する演算が整備されている。
itertools の permutations(n) を用いると順列を生成してくれる。要素の型はベクタ。
.all(f) はイテレータの要素 i が判定関数 f により全て true を返すかどうかを判定する。
- 似た関数に .any(f) があり、これは一つ以上の要素が true を返すかどうかを判定する。
- これを用いると check を一行で求めることもできる。
```
let check = (0..n).permutations(n).any(
    |p| (0..m).all(|i| edge.contains(&(p[cd[i].0], p[cd[i].1]))
));
```

ABC231 C - Counting 2

問題：

Python

提出：https://atcoder.jp/contests/abc231/submissions/31439323

I = lambda:int(input())
LI = lambda:list(map(int,input().split()))
from bisect import bisect_left

n,q = LI()
al = sorted(LI())
query = [I() for _  in [0]*q]
for q in query:
    print(n - bisect_left(al,q))

Rust

提出：https://atcoder.jp/contests/abc231/submissions/31382159

use superslice::Ext;
fn main() {
    proconio::input! {
        n:usize,q:usize,mut al:[usize;n],query:[usize;q]
    }
    al.sort();
    for x in query {
        println!("{}", n - al.lower_bound(&x));
    }
}

二分探索をする。
Rust には superslice に lower_bound と upper_bound がある。

ABC213 C - Reorder Cards

問題：

Python

提出：https://atcoder.jp/contests/abc213/submissions/31439396

I = lambda:int(input())
LI = lambda:list(map(int,input().split()))

h,w,n = LI()
ab = [LI() for _ in [0]*n]
al,bl = zip(*ab)
a2i = {v:i for i,v in enumerate(sorted(set(al)))}
b2i = {v:i for i,v in enumerate(sorted(set(bl)))}
for a,b in ab:
    print(a2i[a] + 1,b2i[b] + 1)

Rust

提出：https://atcoder.jp/contests/abc213/submissions/31420996

use proconio::marker::Usize1;
use std::collections::HashMap;

fn main() {
    proconio::input! {
        _:usize,_:usize,n:usize,
        ab:[(Usize1,Usize1);n]
    }
    let (mut al, mut bl): (Vec<_>, Vec<_>) = ab.iter().cloned().unzip();
    al.sort();
    al.dedup();
    bl.sort();
    bl.dedup();
    let zaatu = |lis: &Vec<_>| {
        lis.iter()
            .enumerate()
            .map(|(i, &a)| (a, i))
            .collect::<HashMap<_, _>>()
    };
    let a2i = zaatu(&al);
    let b2i = zaatu(&bl);
    for (a, b) in ab {
        println!("{} {}", a2i[&a] + 1, b2i[&b] + 1);
    }
}

座標圧縮する。
itr.cloned() は itr.map(|v| v.clone()) と同じ。
- lis.clone().into_iter() でもよい。into_iter() を使う場合は所有権が移動することに注意して使う必要がある。
unzip はタプルのリストをタプルのインデックスごとに複数のベクタに分解する。

Python では重複を除いたソート済配列を sorted(set(al)) で作れるが、Rust では重複削除は al.sort(); と al.dedup(); の組み合わせで実現することが多い。

HashSet を用いて一行で書く方法もある

use std::collections::HashSet;
use itertools::{sorted, Itertools};

fn main() {
    let al = vec![3, 3, 3, 4, 4, 2, 3, 1];
    let bl = sorted(al.clone().into_iter().collect::<HashSet<_>>()).collect_vec();
    println!("{:?}", al); // [3, 3, 3, 4, 4, 2, 3, 1]
    println!("{:?}", bl); // [1, 2, 3, 4]
}