69 lines
1.4 KiB
Rust
69 lines
1.4 KiB
Rust
use std::collections::BTreeSet;
|
|
|
|
struct SmallestInfiniteSet {
|
|
smaller: BTreeSet<i32>,
|
|
next_infinite: i32,
|
|
}
|
|
|
|
/**
|
|
* `&self` means the method takes an immutable reference.
|
|
* If you need a mutable reference, change it to `&mut self` instead.
|
|
*/
|
|
impl SmallestInfiniteSet {
|
|
fn new() -> Self {
|
|
Self {
|
|
smaller: BTreeSet::new(),
|
|
next_infinite: 1,
|
|
}
|
|
}
|
|
|
|
fn pop_smallest(&mut self) -> i32 {
|
|
if let Some(&m) = self.smaller.iter().min() {
|
|
self.smaller.remove(&m);
|
|
|
|
return m;
|
|
}
|
|
|
|
self.next_infinite += 1;
|
|
self.next_infinite - 1
|
|
}
|
|
|
|
fn add_back(&mut self, num: i32) {
|
|
if num >= self.next_infinite {
|
|
return;
|
|
}
|
|
|
|
self.smaller.insert(num);
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Your SmallestInfiniteSet object will be instantiated and called as such:
|
|
* let obj = SmallestInfiniteSet::new();
|
|
* let ret_1: i32 = obj.pop_smallest();
|
|
* obj.add_back(num);
|
|
*/
|
|
|
|
fn main() {}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn example_1() {
|
|
let mut s = SmallestInfiniteSet::new();
|
|
|
|
s.add_back(2);
|
|
|
|
assert_eq!(s.pop_smallest(), 1);
|
|
assert_eq!(s.pop_smallest(), 2);
|
|
assert_eq!(s.pop_smallest(), 3);
|
|
|
|
s.add_back(1);
|
|
|
|
assert_eq!(s.pop_smallest(), 1);
|
|
assert_eq!(s.pop_smallest(), 4);
|
|
assert_eq!(s.pop_smallest(), 5);
|
|
}
|
|
}
|