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

参考