はじめに
AtCoder Beginner Contest 222 の E問題までを Rust で解きました！
詳しめに解説をしていこうと思います。
役にたちましたら LGTM いただけると嬉しいです。
AC 確認はしてますが、もっとスマートな書き方がありましたら教えてください。
なぜやるのか
コンテストから時間をあけて解きなおすことで、記憶へ定着すればいいなと思っています。
また、会社で Rust を使うのですが私はまだ慣れておらず、その練習のためにやっています。
(コンテスト本番には C++ を使って参加している水色コーダーです。)
目次
A - Four Digits (Difficulty 5)
B - Failing Grade (Difficulty 9)
C - Swiss-System Tournament (Difficulty 367)
D - Between Two Arrays (Difficulty 865)
E - Red and Blue Tree (Difficulty 1491)
前提
この記事にあるコードはこちらを import しています。
#[allow(unused_imports)]
use proconio::{input, fastout, marker::Chars};
#[allow(unused_imports)]
use std::collections::{HashSet, HashMap, BTreeSet, VecDeque};
#[allow(unused_imports)]
use std::cmp::{max, min};
#[allow(unused_imports)]
use itertools::Itertools;
A - Four Digits
$4 - |N|$ 文字だけ 0 を出力した後、 $N$ を出力します。(全体で複数回出力します。)
$N$ の先頭に 0 を足して、まとめて出力する方法もあります。
fn main() {
input! {
n: Chars
}
for _ in 0..4 - n.len() {
print!("0");
}
print!("{}", n.iter().collect::<String>());
}
B - Failing Grade
$P$ 未満の $a_i$ の個数が答えになりますね。
iter() や into_iter() などからつなげられる filter() と count() を使うと簡単に書けます。
fn main() {
input! {
n: usize,
p: usize,
a: [usize; n]
}
let ans = a.into_iter().filter(|v| v < &p).count();
println!("{}", ans);
}
C - Swiss-System Tournament
問題はわかるのですが、実装が複雑ですね。
まず実装はおいておいて、問題にあるじゃんけん大会の流れからみます。
以下の流れになります。
- 順位で対戦相手を決める (全Nグループ)
- じゃんけんをする
- 勝数から順位をつける
- 次のラウンドに行く
次に、それぞれの処理を考えます。
- について、
2kと2k + 1番目の人が対戦相手になります。 - について、じゃんけんで何を出すかは入力でもらえますね。勝敗をきめる部分は全通り調べます。
- について、順位を決めるのは勝数と番号ですね。勝敗の結果と、はじめのインデックスを使ってソートします。
- について、1~3 を
M回 くり返します。
3 について、勝数は降順・番号は昇順でソートするので注意が必要です。
私は、 (勝数 * -1, 番号) という tuple の Vec をつかいました。こうすることで、 Vec.sort() だけで順位をつけられます。
// c1 と c2 のじゃんけんの結果
// 0 はひきわけ
// 1 は c1 の勝ち
// 2 は c2 の勝ち
fn battle (c1: char, c2: char) -> usize {
if c1 == c2 { return 0 }
else if c1 == 'G' && c2 == 'C' { return 1 }
else if c1 == 'G' && c2 == 'P' { return 2 }
else if c1 == 'C' && c2 == 'G' { return 2 }
else if c1 == 'C' && c2 == 'P' { return 1 }
else if c1 == 'P' && c2 == 'G' { return 1 }
else if c1 == 'P' && c2 == 'C' { return 2 }
100 // ここにはこない
}
fn main() {
input! {
n: usize,
m: usize,
a: [Chars; 2 * n]
}
let mut score = vec![(0_i32, 0_i32); 2 * n];
for i in 0..2 * n {
score[i].1 = i as i32; // 番号をつける
}
for i in 0..m { // 4. M回くりかえす
for j in 0..n { // 1. Nグループでじゃんけんをする
let p1 = score[j * 2].1 as usize; // 1. 対戦相手をきめる
let p2 = score[j * 2 + 1].1 as usize; // 1. 対戦相手をきめる
let result = battle(a[p1][i], a[p2][i]); // 2. じゃんけんをする
if result == 1 { score[j * 2].0 -= 1 } // 3. 順位をつける / 勝ち数は -1 をかけるので -1 をしている
if result == 2 { score[j * 2 + 1].0 -= 1 } // 3. 順位をつける
}
score.sort(); // 3. 順位をつける
}
for i in 0..n * 2 {
println!("{}", score[i].1 + 1);
}
}
D - Between Two Arrays
DP の問題ですね。
遷移を考えるフェーズ1 と、累積和で高速化するフェーズ2 があります。
まずはフェーズ1 について
$c_{i+1}$ が $M$ の時の個数は、$c_i$ が $M$ 以下の時の個数の総和になりますね。
この遷移をすれば解けるのですが、愚直に実装すると $O(N M^2)$ となり TLE です。
フェーズ2 として高速化をします。
求める値は $0 \sim M$ の和なので、累積和が使えそうです。
累積和の計算量は前処理が $O(M)$ 、区間和の計算が $O(1)$ です。そのため、全体で $O(N M)$ になって時間内に解けました。
続いて実装ですが注意が必要で、$c_i$ の範囲が $a_i \le c_i \le b_i$ なので、$c_i$ から $c_{i+1}$ を計算するときは $a_i \sim b_i$ の範囲の値のみを使うようにします。また、以降の遷移で使うので、累積和は $c_{i+1}$ の値に関係なく $[0, 3000]$ の範囲を計算しています。
fn main() {
input! {
n: usize,
a: [usize; n],
b: [usize; n]
}
let m = 998244353;
// dp[i] = c の値が i の時の組み合わせ
let mut dp = vec![1; 3010];
for i in 0..n {
let mut next_dp = vec![0; 3010];
for c in a[i]..=b[i] { // a_i ~ b_i が計算対象になる
next_dp[c] = dp[c]
}
for v in 0..=3000 { // 累積和を求める
next_dp[v + 1] += next_dp[v];
next_dp[v + 1] %= m;
}
dp = next_dp;
}
println!("{}", dp[b[n - 1]]);
}
E - Red and Blue Tree
グラフとDPが組み合わさった問題です。
まず、それぞれの枝を通る回数を求めます。回数を求めるには、少し難しいですが経路の復元が必要です。
経路を復元するためには、初めての頂点に訪れるたびに直前の枝のインデックスと直前のノードを記録します。目的地についたら、スタート地点まで逆向きにたどっていくことで経路が復元できます。目的地にいくのが $O(N)$ で、スタート地点に戻るのが $O(N)$ なので、全体で $O(N)$ です。
ちなみに ABC218-F でも似た処理が出ました。
これを $M$ 回やるので、この部分の計算量は $O(NM)$ になります。
次はDPです。
dp[i] = 和が i になるときの組み合わせ という dpテーブルで解けます。これはよくある DP で、 ABC216-F でも出ました。
ちなみに、dp[i][j] = 枝 i まで見て和が j になるときの組み合わせ とするとメモリ制限超過になるので注意です。
DPテーブルはできましたが、 k は負の値もとるので困りますね。こういうときは、 dp[k] を「和が 0 のとき」として扱います。
dp[0] が 和が「-k のとき」を示して、 dp[2*k] が「和が k のとき」を示すようにします。
こうすると、配列のインデックスが負にならずに済みます。
また、K の代わりに、取り得る最大値 (コードで言う use_count.iter().sum(); ) を使うことで、少しばかりのメモリ節約と高速化を狙っています。
DP部分の計算量は $O(N |K|)$ になります。
全体の計算量は $O(NM + N |K|)$ となりました。
fn main() {
input! {
n: usize,
m: usize,
k: i64,
a: [usize; m],
uv: [(usize, usize); n - 1]
}
let mut graph = vec![vec![]; n];
for i in 0..n - 1 {
let (u, v) = uv[i];
let (u, v) = (u - 1, v - 1);
graph[u].push((v, i));
graph[v].push((u, i));
}
let inf = 1001001001;
let mut use_count = vec![0; n - 1]; // 枝を使った回数
// from から to までの最短経路で通る枝を記録する
let mut bfs_shortest_path_index = |from: usize, to: usize| {
// prev[i] = (頂点, index)
let mut prev = vec![(inf, inf); n];
let mut dist = vec![inf; n];
dist[from] = 0;
let mut que = VecDeque::new();
que.push_back(from);
// BFS で最短経路を求める
while let Some(now) = que.pop_front() {
let depth = dist[now];
if now == to { break }
for &(next, idx) in graph[now].iter() {
if depth + 1 < dist[next] {
dist[next] = depth + 1;
prev[next] = (now, idx);
que.push_back(next);
}
}
}
// 使った枝を記録する
let mut now = to;
while now != from {
use_count[prev[now].1] += 1;
now = prev[now].0;
}
};
for i in 1..m {
bfs_shortest_path_index(a[i - 1] - 1, a[i] - 1);
}
let m = 998244353_usize;
use_count.sort(); // ソートして小さい順に処理する
let sum: usize = use_count.iter().sum();
let mut sumv = 0;
if max(k, -k) as usize > sum {
println!("{}", 0);
return;
}
// dp[i] = 和が i になるときの組み合わせ
let mut dp = vec![0_usize; sum * 2 + 1_usize];
dp[sum] = 1; // dp[sum] が 和が 1 の時の組み合わせを示す
for i in 0..n - 1 {
let mut next_dp = vec![0_usize; sum * 2 + 1_usize];
sumv += use_count[i]; // これまでの累積和
for j in sum - sumv..=sum + sumv { // ソートして小さい順に処理することで、dpテーブルの更新範囲を小さくする
if j + use_count[i] <= sum * 2 {
next_dp[j] += dp[j + use_count[i]];
next_dp[j] %= m;
}
if j >= use_count[i] {
next_dp[j] += dp[j - use_count[i]];
next_dp[j] %= m;
}
}
dp = next_dp;
}
println!("{}", dp[(sum as i64 + k) as usize]);
}
最後まで読んでいただきありがとうございます！