Heap / Priority Queue - Smallest Number in Infinite Set

Design a class that represents a set containing all positive integers. Support popSmallest() and addBack(num) operations.

popSmallest()→1, popSmallest()→2, addBack(1), popSmallest()→1, popSmallest()→3

Use a pointer for the smallest "unpopped" consecutive integer, and a min-heap for numbers that were added back.

Maintain cur (next consecutive unpoped, starts at 1) and a min-heap for added-back numbers.
addedBack Set to avoid duplicates in heap.
popSmallest: return min(heap.top, cur). If heap non-empty and heap.top < cur, pop heap. Else return cur++.
addBack(num): if num < cur and not in addedBack, add to heap and set.

import java.util.*; class SmallestInfiniteSet { PriorityQueue<Integer> heap = new PriorityQueue<>(); Set<Integer> added = new HashSet<>(); int cur = 1; public int popSmallest() { if (!heap.isEmpty() && heap.peek() < cur) { int val = heap.poll(); added.remove(val); return val; } return cur++; } public void addBack(int num) { if (num < cur && !added.contains(num)) { heap.offer(num); added.add(num); } } }

Time Complexity: pop: O(log N), addBack: O(log N)
Space Complexity: O(N)