LoginSignup
0
0

More than 3 years have passed since last update.

Rustでmultisetもどき(ABC170 - E)

Posted at

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触ってみて数日なので大変だったけど動いてよかった。おしまい。

参考

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0