LeetCode/rs/smallest-number-in-infinite-set.rs

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);
    }
}