AtCoderの E - Smart Infants をRustで解いてみました。
C++でやったときは(解説を読んで)multisetを使ったのですが同等のものがRust(のstd?)には無いっぽい(自分調べ)ので同じ働きをする構造体MultiSetを作りました。
↑の問題が解ければ嬉しいのでその場しのぎ的な実装ですが。
できること
- ひとまず
usize
のみ対応 - 重複許可して挿入
- 要素を1つ削除
- 最小値の取得
- 最大値の取得
概要
-
std::collections::BTreeSet
で値を保存 - 重複している個数を
std::collections::HashMap
で管理
実装
use std::collections::{BTreeSet, HashMap};
#[derive(Clone, Debug)]
struct MultiSet {
set: BTreeSet<usize>,
map: HashMap<usize, usize>
}
impl MultiSet {
fn new() -> MultiSet {
MultiSet{
set: BTreeSet::new(), map: HashMap::new()
}
}
///multiset的に書き出し
fn print(&self)
{
print!("{{ ");
for x in &self.set {
if let Some(&num) = self.map.get(x){
for _i in 0..num {
print!("{} ", x);
}
}
}
println!("}}");
}
///重複許可挿入
fn insert(&mut self, i:usize) -> Option<usize>
{
if let Some(_i) = self.set.get(&i) {
//setにある
*self.map.entry(i).or_insert(0) += 1;
}else{
//setにない
self.set.insert(i);
*self.map.entry(i).or_insert(0) += 1;
}
return Some(i);
}
///1つ削除
fn erase(&mut self, e:usize) -> Option<usize>
{
if let Some(_e) = self.set.get(&e) {
//setにある
*self.map.entry(e).or_insert(0) -= 1;
if self.map[&e]==0 {
//なくなった
self.set.take(&e);
}
return Some(e);
}else{
//setにない
return None;
}
}
///最小値の取得
fn get_min(&self) -> Option<usize>
{
if let Some(&m) = self.set.iter().nth(0) {
return Some(m);
}else{
return None;
}
}
///最大値の取得
fn get_max(&self) -> Option<usize>
{
if let Some(&m) = self.set.iter().last() {
return Some(m);
}else{
return None;
}
}
}
使用例
書き出し
///multiset的に書き出し
fn print(&self)
{
print!("{{ ");
for x in &self.set {
if let Some(&num) = self.map.get(x){
for _i in 0..num {
print!("{} ", x);
}
}
}
println!("}}");
}
あたかも『multiset』であるかのように書き出す。set
(昇順に整理されている)の中身を見ながら、map
に記録してある数だけ書き出す。
挿入
//略
///重複許可挿入
fn insert(&mut self, i:usize) -> Option<usize>
{
if let Some(_i) = self.set.get(&i) {
//setにある
*self.map.entry(i).or_insert(0) += 1;
}else{
//setにない
self.set.insert(i);
*self.map.entry(i).or_insert(0) += 1;
}
return Some(i);
}
//略
fn main(){
let mut myset = MultiSet::new();
myset.insert(3);
myset.insert(5);
myset.insert(3);
myset.insert(6);
myset.print();
println!("{:?}", myset);
}
{ 3 3 5 6 }
MultiSet { set: {3, 5, 6}, map: {6: 1, 3: 2, 5: 1} }
#[derive(Clone, Debug)]
でDebugをつけてるので内部の状態がわかる(ざっくり)
削除
//略
///1つ削除
fn erase(&mut self, e:usize) -> Option<usize>
{
if let Some(_e) = self.set.get(&e) {
//setにある
*self.map.entry(e).or_insert(0) -= 1;
if self.map[&e]==0 {
//なくなった
self.set.take(&e);
}
return Some(e);
}else{
//setにない
return None;
}
}
//略
fn main(){
//上の続きで
myset.erase(3);
assert_eq!( myset.erase(4), None);
myset.erase(6);
myset.print();
println!("{:?}", myset);
}
{ 3 5 }
MultiSet { set: {3, 5}, map: {6: 0, 3: 1, 5: 1} }
set
にないものは削除できないのでNone
を返す。削除して数が0になったらset
から削除している。
最大最小
//略
///最小値の取得
fn get_min(&self) -> Option<usize>
{
if let Some(&m) = self.set.iter().nth(0) {
return Some(m);
}else{
return None;
}
}
///最大値の取得
fn get_max(&self) -> Option<usize>
{
if let Some(&m) = self.set.iter().last() {
return Some(m);
}else{
return None;
}
}
//略
fn main(){
//さらに続きで
myset.insert(9);
myset.insert(1);
myset.print();
if let Some(min) = myset.get_min() {println!("min = {}", min)}
if let Some(max) = myset.get_max() {println!("max = {}", max)}
println!("{:?}", myset);
}
{ 1 3 5 9 }
min = 1
max = 9
MultiSet { set: {1, 3, 5, 9}, map: {1: 1, 9: 1, 5: 1, 3: 1, 6: 0} }
わざわざ中で分岐させなくてもよかったのでは?
AtCoderに提出
AC → https://atcoder.jp/contests/abc170/submissions/15809165
おわり
真面目にRust触ってみて数日なので大変だったけど動いてよかった。おしまい。