problems(rs): add “2336. Smallest Number in Infinite Set”

Signed-off-by: Matej Focko <mfocko@redhat.com>

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

@@ -0,0 +1,69 @@
+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);
+    }
+}