use std::collections::BTreeSet; struct SmallestInfiniteSet { smaller: BTreeSet, 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); } }