2336. Smallest Number in Infinite Set

MediumDesignHash TableOrdered SetHeap (Priority Queue)

Problem Description

You need to implement a data structure that represents an infinite set containing all positive integers [1, 2, 3, 4, 5, ...].

The SmallestInfiniteSet class should support the following operations:

Constructor SmallestInfiniteSet(): Initializes the object to contain all positive integers starting from 1.
Pop Smallest popSmallest(): Removes and returns the smallest integer currently in the set. For example, if the set contains [1, 2, 3, ...], calling this method would remove and return 1, leaving the set as [2, 3, 4, ...].
Add Back addBack(num): Adds a positive integer num back into the set, but only if it's not already present. This is useful when you want to restore a number that was previously popped.

The key challenge is efficiently managing which numbers have been removed from the infinite set and which have been added back. While the set conceptually contains all positive integers, you need to track the state changes caused by the pop and add operations.

For example:

Initially, the set contains [1, 2, 3, 4, 5, ...]
After calling popSmallest(), it returns 1 and the set becomes [2, 3, 4, 5, ...]
After calling popSmallest() again, it returns 2 and the set becomes [3, 4, 5, ...]
After calling addBack(1), the set becomes [1, 3, 4, 5, ...]
The next popSmallest() would return 1

Quick Interview Experience

Help others by sharing your interview experience

Have you seen this problem before?

Intuition

Since we need to track an "infinite" set of positive integers with operations to remove the smallest element and add elements back, we need to think about what actually changes in our set.

Initially, the set contains all positive integers [1, 2, 3, 4, ...]. When we pop elements, we create "gaps" in this sequence. For instance, if we pop 1 and 2, our set becomes [3, 4, 5, ...]. If we then add back 1, it becomes [1, 3, 4, 5, ...].

The key insight is that we don't actually need to store an infinite set. We only need to track which numbers have been removed from the initial sequence and haven't been added back yet. Since the problem constraints typically limit the number of operations (and based on the solution using range(1, 1001)), we can work with a finite range.

We need a data structure that:

Maintains elements in sorted order (to quickly find the smallest)
Allows efficient removal of the minimum element
Allows efficient insertion of elements

An ordered set (like SortedSet in Python) is perfect for this because:

It automatically maintains elements in sorted order
The first element s[0] is always the smallest
Both insertion and removal operations are O(log n)

Instead of conceptually thinking about tracking removed elements, we can flip our perspective: maintain a set of all available numbers. Start with a finite range like [1, 1000] (sufficient for the problem constraints), and when popSmallest() is called, remove the smallest element from our set. When addBack(num) is called, add the number back to our set if it was previously removed.

This approach elegantly handles all cases - the set maintains order automatically, duplicate prevention is built-in (sets don't allow duplicates), and we always have immediate access to the smallest element.

Learn more about Heap (Priority Queue) patterns.

Solution Approach

The solution uses an ordered set data structure to efficiently manage the infinite set operations. Here's how each component works:

Data Structure Choice

We use a SortedSet (an ordered set implementation) initialized with elements from 1 to 1000. This range is sufficient based on the problem constraints. The ordered set maintains elements in sorted order automatically, which is crucial for finding the smallest element quickly.

self.s = SortedSet(range(1, 1001))

Implementation Details

1. Initialization (__init__)

Create a SortedSet containing all integers from 1 to 1000
Time Complexity: O(n × log n) where n = 1000
The set starts with all possible values we might need

2. Pop Smallest (popSmallest)

Access the first element using index s[0] - this is always the smallest due to the sorted property
Remove this element from the set using remove()
Return the removed element
Time Complexity: O(log n) for the removal operation

def popSmallest(self) -> int:
    x = self.s[0]      # Get the smallest element
    self.s.remove(x)   # Remove it from the set
    return x           # Return the removed element

3. Add Back (addBack)

Simply add the number back to the set using add()
The SortedSet automatically:
- Maintains the sorted order
- Handles duplicates (won't add if already present)
Time Complexity: O(log n) for the insertion operation

def addBack(self, num: int) -> None:
    self.s.add(num)    # Add the number back to the set

Why This Works

The elegance of this solution lies in treating the "infinite set" as a finite ordered collection of available numbers. Since we're limited by the problem constraints (numbers up to 1000), we can:

Pre-populate all possible values
Remove numbers when popped
Add them back when requested

The SortedSet handles all the complexity of maintaining order and ensuring efficient operations, making the implementation clean and straightforward.

Ready to land your dream job?

Unlock your dream job with a 5-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through a concrete example to see how the solution works:

Initial State:

We create a SmallestInfiniteSet
The SortedSet contains: [1, 2, 3, 4, 5, ..., 1000]

Operation 1: popSmallest()

Current set: [1, 2, 3, 4, 5, ...]
Access s[0] which is 1 (the smallest element)
Remove 1 from the set
Return 1
New set: [2, 3, 4, 5, 6, ...]

Operation 2: popSmallest()

Current set: [2, 3, 4, 5, 6, ...]
Access s[0] which is now 2 (the new smallest)
Remove 2 from the set
Return 2
New set: [3, 4, 5, 6, 7, ...]

Operation 3: popSmallest()

Current set: [3, 4, 5, 6, 7, ...]
Access s[0] which is 3
Remove 3 from the set
Return 3
New set: [4, 5, 6, 7, 8, ...]

Operation 4: addBack(1)

Current set: [4, 5, 6, 7, 8, ...]
Add 1 to the set
The SortedSet automatically places it in the correct position
New set: [1, 4, 5, 6, 7, ...]

Operation 5: popSmallest()

Current set: [1, 4, 5, 6, 7, ...]
Access s[0] which is 1 (it's back as the smallest!)
Remove 1 from the set
Return 1
New set: [4, 5, 6, 7, 8, ...]

Operation 6: addBack(2)

Current set: [4, 5, 6, 7, 8, ...]
Add 2 to the set
The SortedSet automatically maintains sorted order
New set: [2, 4, 5, 6, 7, ...]

Operation 7: addBack(2) (duplicate attempt)

Current set: [2, 4, 5, 6, 7, ...]
Try to add 2 to the set
Since 2 is already present, the set remains unchanged
New set: [2, 4, 5, 6, 7, ...] (no change)

This example demonstrates how the SortedSet elegantly handles:

Always maintaining sorted order (smallest element is always at index 0)
Efficient removal and insertion operations
Automatic duplicate prevention
The illusion of an "infinite set" while actually working with a finite range

Solution Implementation

1from sortedcontainers import SortedSet
2
3class SmallestInfiniteSet:
4    def __init__(self):
5        # Initialize a sorted set containing integers from 1 to 1000
6        # This represents the "infinite" set of positive integers
7        self.sorted_set = SortedSet(range(1, 1001))
8
9    def popSmallest(self) -> int:
10        # Get the smallest element (first element in sorted set)
11        smallest = self.sorted_set[0]
12      
13        # Remove the smallest element from the set
14        self.sorted_set.remove(smallest)
15      
16        # Return the smallest element that was removed
17        return smallest
18
19    def addBack(self, num: int) -> None:
20        # Add the number back to the set
21        # SortedSet automatically handles duplicates (won't add if already exists)
22        self.sorted_set.add(num)
23
24
25# Your SmallestInfiniteSet object will be instantiated and called as such:
26# obj = SmallestInfiniteSet()
27# param_1 = obj.popSmallest()
28# obj.addBack(num)
29

1class SmallestInfiniteSet {
2    // TreeSet to maintain sorted unique integers in ascending order
3    private TreeSet<Integer> availableNumbers = new TreeSet<>();
4
5    /**
6     * Constructor initializes the set with integers from 1 to 1000
7     * representing the smallest positive integers in the infinite set
8     */
9    public SmallestInfiniteSet() {
10        // Populate the set with initial values 1 through 1000
11        for (int i = 1; i <= 1000; i++) {
12            availableNumbers.add(i);
13        }
14    }
15
16    /**
17     * Removes and returns the smallest integer from the infinite set
18     * @return the smallest available integer
19     */
20    public int popSmallest() {
21        // Remove and return the first (smallest) element from the TreeSet
22        return availableNumbers.pollFirst();
23    }
24
25    /**
26     * Adds a number back to the infinite set if it's not already present
27     * @param num the integer to add back to the set
28     */
29    public void addBack(int num) {
30        // Add the number back to the set (TreeSet handles duplicates automatically)
31        availableNumbers.add(num);
32    }
33}
34
35/**
36 * Your SmallestInfiniteSet object will be instantiated and called as such:
37 * SmallestInfiniteSet obj = new SmallestInfiniteSet();
38 * int param_1 = obj.popSmallest();
39 * obj.addBack(num);
40 */
41

1class SmallestInfiniteSet {
2public:
3    // Constructor: Initialize the set with integers from 1 to 1000
4    SmallestInfiniteSet() {
5        for (int i = 1; i <= 1000; ++i) {
6            available_numbers.insert(i);
7        }
8    }
9  
10    // Remove and return the smallest integer from the infinite set
11    int popSmallest() {
12        // Get the smallest element (first element in ordered set)
13        int smallest = *available_numbers.begin();
14      
15        // Remove the smallest element from the set
16        available_numbers.erase(available_numbers.begin());
17      
18        return smallest;
19    }
20  
21    // Add a number back to the infinite set if not already present
22    void addBack(int num) {
23        // Insert the number (set automatically handles duplicates)
24        available_numbers.insert(num);
25    }
26  
27private:
28    // Ordered set to maintain available numbers in sorted order
29    std::set<int> available_numbers;
30};
31
32/**
33 * Your SmallestInfiniteSet object will be instantiated and called as such:
34 * SmallestInfiniteSet* obj = new SmallestInfiniteSet();
35 * int param_1 = obj->popSmallest();
36 * obj->addBack(num);
37 */
38

1// Type alias for comparison function
2type CompareFunction<T> = (left: T, right: T) => number;
3
4// Global TreeSet instance to store the infinite set
5let infiniteSet: TreeSet<number>;
6
7// Initialize the smallest infinite set with numbers 1 to 1000
8function initializeSmallestInfiniteSet(): void {
9    infiniteSet = new TreeSet<number>();
10    for (let i = 1; i <= 1000; i++) {
11        infiniteSet.add(i);
12    }
13}
14
15// Remove and return the smallest integer from the infinite set
16function popSmallest(): number {
17    return infiniteSet.shift()!;
18}
19
20// Add a number back to the infinite set if it's not already present
21function addBack(num: number): void {
22    infiniteSet.add(num);
23}
24
25// Red-Black Tree Node class
26class RBTreeNode<T = number> {
27    data: T;
28    count: number; // Number of duplicates for this value
29    left: RBTreeNode<T> | null;
30    right: RBTreeNode<T> | null;
31    parent: RBTreeNode<T> | null;
32    color: number; // 0 = RED, 1 = BLACK
33
34    constructor(data: T) {
35        this.data = data;
36        this.left = null;
37        this.right = null;
38        this.parent = null;
39        this.color = 0; // New nodes are RED by default
40        this.count = 1;
41    }
42
43    // Get the sibling of this node
44    sibling(): RBTreeNode<T> | null {
45        if (!this.parent) return null;
46        return this.isOnLeft() ? this.parent.right : this.parent.left;
47    }
48
49    // Check if this node is the left child of its parent
50    isOnLeft(): boolean {
51        return this === this.parent!.left;
52    }
53
54    // Check if this node has at least one red child
55    hasRedChild(): boolean {
56        return (
57            Boolean(this.left && this.left.color === 0) ||
58            Boolean(this.right && this.right.color === 0)
59        );
60    }
61}
62
63// Red-Black Tree implementation
64class RBTree<T> {
65    root: RBTreeNode<T> | null;
66    lessThan: (left: T, right: T) => boolean;
67
68    constructor(compare: CompareFunction<T> = (left: T, right: T) => (left < right ? -1 : left > right ? 1 : 0)) {
69        this.root = null;
70        this.lessThan = (left: T, right: T) => compare(left, right) < 0;
71    }
72
73    // Perform left rotation on the given node
74    rotateLeft(pivotNode: RBTreeNode<T>): void {
75        const rightChild = pivotNode.right!;
76        pivotNode.right = rightChild.left;
77
78        if (pivotNode.right) {
79            pivotNode.right.parent = pivotNode;
80        }
81      
82        rightChild.parent = pivotNode.parent;
83
84        if (!pivotNode.parent) {
85            this.root = rightChild;
86        } else if (pivotNode === pivotNode.parent.left) {
87            pivotNode.parent.left = rightChild;
88        } else {
89            pivotNode.parent.right = rightChild;
90        }
91
92        rightChild.left = pivotNode;
93        pivotNode.parent = rightChild;
94    }
95
96    // Perform right rotation on the given node
97    rotateRight(pivotNode: RBTreeNode<T>): void {
98        const leftChild = pivotNode.left!;
99        pivotNode.left = leftChild.right;
100
101        if (pivotNode.left) {
102            pivotNode.left.parent = pivotNode;
103        }
104      
105        leftChild.parent = pivotNode.parent;
106
107        if (!pivotNode.parent) {
108            this.root = leftChild;
109        } else if (pivotNode === pivotNode.parent.left) {
110            pivotNode.parent.left = leftChild;
111        } else {
112            pivotNode.parent.right = leftChild;
113        }
114
115        leftChild.right = pivotNode;
116        pivotNode.parent = leftChild;
117    }
118
119    // Swap colors between two nodes
120    swapColor(node1: RBTreeNode<T>, node2: RBTreeNode<T>): void {
121        const tempColor = node1.color;
122        node1.color = node2.color;
123        node2.color = tempColor;
124    }
125
126    // Swap data between two nodes
127    swapData(node1: RBTreeNode<T>, node2: RBTreeNode<T>): void {
128        const tempData = node1.data;
129        node1.data = node2.data;
130        node2.data = tempData;
131    }
132
133    // Fix Red-Black Tree properties after insertion
134    fixAfterInsert(currentNode: RBTreeNode<T>): void {
135        let parentNode = null;
136        let grandParentNode = null;
137
138        // Continue fixing until we reach root or find a valid configuration
139        while (currentNode !== this.root && currentNode.color !== 1 && currentNode.parent?.color === 0) {
140            parentNode = currentNode.parent;
141            grandParentNode = currentNode.parent.parent;
142
143            // Case A: Parent is left child of grandparent
144            if (parentNode === grandParentNode?.left) {
145                const uncleNode = grandParentNode.right;
146
147                // Case 1: Uncle is RED - only recoloring needed
148                if (uncleNode && uncleNode.color === 0) {
149                    grandParentNode.color = 0;
150                    parentNode.color = 1;
151                    uncleNode.color = 1;
152                    currentNode = grandParentNode;
153                } else {
154                    // Case 2: Current node is right child - left rotation needed
155                    if (currentNode === parentNode.right) {
156                        this.rotateLeft(parentNode);
157                        currentNode = parentNode;
158                        parentNode = currentNode.parent;
159                    }
160
161                    // Case 3: Current node is left child - right rotation needed
162                    this.rotateRight(grandParentNode);
163                    this.swapColor(parentNode!, grandParentNode);
164                    currentNode = parentNode!;
165                }
166            } else {
167                // Case B: Parent is right child of grandparent
168                const uncleNode = grandParentNode!.left;
169
170                // Case 1: Uncle is RED - only recoloring needed
171                if (uncleNode != null && uncleNode.color === 0) {
172                    grandParentNode!.color = 0;
173                    parentNode.color = 1;
174                    uncleNode.color = 1;
175                    currentNode = grandParentNode!;
176                } else {
177                    // Case 2: Current node is left child - right rotation needed
178                    if (currentNode === parentNode.left) {
179                        this.rotateRight(parentNode);
180                        currentNode = parentNode;
181                        parentNode = currentNode.parent;
182                    }
183
184                    // Case 3: Current node is right child - left rotation needed
185                    this.rotateLeft(grandParentNode!);
186                    this.swapColor(parentNode!, grandParentNode!);
187                    currentNode = parentNode!;
188                }
189            }
190        }
191      
192        // Root must always be BLACK
193        this.root!.color = 1;
194    }
195
196    // Delete a single occurrence of the value
197    delete(value: T): boolean {
198        const nodeToDelete = this.find(value);
199        if (!nodeToDelete) return false;
200      
201        nodeToDelete.count--;
202        if (!nodeToDelete.count) {
203            this.deleteNode(nodeToDelete);
204        }
205        return true;
206    }
207
208    // Delete all occurrences of the value
209    deleteAll(value: T): boolean {
210        const nodeToDelete = this.find(value);
211        if (!nodeToDelete) return false;
212      
213        this.deleteNode(nodeToDelete);
214        return true;
215    }
216
217    // Delete a node from the tree
218    deleteNode(nodeToDelete: RBTreeNode<T>): void {
219        // Find replacement node using BST delete rules
220        const replacementNode = this.findBSTReplacement(nodeToDelete);
221
222        // Check if both nodes are BLACK
223        const bothBlack = (replacementNode === null || replacementNode.color === 1) && nodeToDelete.color === 1;
224        const parentNode = nodeToDelete.parent!;
225
226        if (!replacementNode) {
227            // Node to delete is a leaf
228            if (nodeToDelete === this.root) {
229                this.root = null;
230            } else {
231                if (bothBlack) {
232                    // Both nodes are BLACK - fix double black
233                    this.fixDoubleBlack(nodeToDelete);
234                } else {
235                    // One node is RED - make sibling RED if exists
236                    if (nodeToDelete.sibling()) {
237                        nodeToDelete.sibling()!.color = 0;
238                    }
239                }
240              
241                // Remove node from tree
242                if (nodeToDelete.isOnLeft()) {
243                    parentNode.left = null;
244                } else {
245                    parentNode.right = null;
246                }
247            }
248            return;
249        }
250
251        if (!nodeToDelete.left || !nodeToDelete.right) {
252            // Node has one child
253            if (nodeToDelete === this.root) {
254                // Replace root with child
255                nodeToDelete.data = replacementNode.data;
256                nodeToDelete.left = null;
257                nodeToDelete.right = null;
258            } else {
259                // Replace node with child
260                if (nodeToDelete.isOnLeft()) {
261                    parentNode.left = replacementNode;
262                } else {
263                    parentNode.right = replacementNode;
264                }
265                replacementNode.parent = parentNode;
266              
267                if (bothBlack) {
268                    this.fixDoubleBlack(replacementNode);
269                } else {
270                    replacementNode.color = 1;
271                }
272            }
273            return;
274        }
275
276        // Node has two children - swap with successor and delete
277        this.swapData(replacementNode, nodeToDelete);
278        this.deleteNode(replacementNode);
279    }
280
281    // Find replacement node for BST deletion
282    private findBSTReplacement(node: RBTreeNode<T>): RBTreeNode<T> | null {
283        // Two children - find inorder successor
284        if (node.left && node.right) {
285            return this.findSuccessor(node.right);
286        }
287      
288        // Leaf node
289        if (!node.left && !node.right) {
290            return null;
291        }
292      
293        // Single child
294        return node.left ?? node.right;
295    }
296
297    // Find the inorder successor (leftmost node in right subtree)
298    private findSuccessor(node: RBTreeNode<T>): RBTreeNode<T> {
299        let current = node;
300        while (current.left) {
301            current = current.left;
302        }
303        return current;
304    }
305
306    // Fix double black violation in Red-Black Tree
307    fixDoubleBlack(node: RBTreeNode<T>): void {
308        // Base case: reached root
309        if (node === this.root) return;
310
311        const siblingNode = node.sibling();
312        const parentNode = node.parent!;
313      
314        if (!siblingNode) {
315            // No sibling - push double black up
316            this.fixDoubleBlack(parentNode);
317        } else {
318            if (siblingNode.color === 0) {
319                // Sibling is RED
320                parentNode.color = 0;
321                siblingNode.color = 1;
322              
323                if (siblingNode.isOnLeft()) {
324                    this.rotateRight(parentNode);
325                } else {
326                    this.rotateLeft(parentNode);
327                }
328              
329                this.fixDoubleBlack(node);
330            } else {
331                // Sibling is BLACK
332                if (siblingNode.hasRedChild()) {
333                    // Sibling has at least one RED child
334                    if (siblingNode.left && siblingNode.left.color === 0) {
335                        if (siblingNode.isOnLeft()) {
336                            // Left-Left case
337                            siblingNode.left.color = siblingNode.color;
338                            siblingNode.color = parentNode.color;
339                            this.rotateRight(parentNode);
340                        } else {
341                            // Right-Left case
342                            siblingNode.left.color = parentNode.color;
343                            this.rotateRight(siblingNode);
344                            this.rotateLeft(parentNode);
345                        }
346                    } else {
347                        if (siblingNode.isOnLeft()) {
348                            // Left-Right case
349                            siblingNode.right!.color = parentNode.color;
350                            this.rotateLeft(siblingNode);
351                            this.rotateRight(parentNode);
352                        } else {
353                            // Right-Right case
354                            siblingNode.right!.color = siblingNode.color;
355                            siblingNode.color = parentNode.color;
356                            this.rotateLeft(parentNode);
357                        }
358                    }
359                    parentNode.color = 1;
360                } else {
361                    // Sibling has two BLACK children
362                    siblingNode.color = 0;
363                  
364                    if (parentNode.color === 1) {
365                        this.fixDoubleBlack(parentNode);
366                    } else {
367                        parentNode.color = 1;
368                    }
369                }
370            }
371        }
372    }
373
374    // Insert a new value into the tree
375    insert(data: T): boolean {
376        // Find insertion position
377        let parentNode = this.root;
378        while (parentNode) {
379            if (this.lessThan(data, parentNode.data)) {
380                if (!parentNode.left) break;
381                parentNode = parentNode.left;
382            } else if (this.lessThan(parentNode.data, data)) {
383                if (!parentNode.right) break;
384                parentNode = parentNode.right;
385            } else {
386                // Value already exists - increment count
387                parentNode.count++;
388                return false;
389            }
390        }
391
392        // Create and insert new node
393        const newNode = new RBTreeNode(data);
394      
395        if (!parentNode) {
396            this.root = newNode;
397        } else if (this.lessThan(newNode.data, parentNode.data)) {
398            parentNode.left = newNode;
399        } else if (this.lessThan(parentNode.data, newNode.data)) {
400            parentNode.right = newNode;
401        } else {
402            parentNode.count++;
403            return false;
404        }
405      
406        newNode.parent = parentNode;
407        this.fixAfterInsert(newNode);
408        return true;
409    }
410
411    // Find a node with the given value
412    find(data: T): RBTreeNode<T> | null {
413        let currentNode = this.root;
414      
415        while (currentNode) {
416            if (this.lessThan(data, currentNode.data)) {
417                currentNode = currentNode.left;
418            } else if (this.lessThan(currentNode.data, data)) {
419                currentNode = currentNode.right;
420            } else {
421                return currentNode;
422            }
423        }
424      
425        return null;
426    }
427
428    // Generator for in-order traversal
429    *inOrder(rootNode: RBTreeNode<T> = this.root!): Generator<T, undefined, void> {
430        if (!rootNode) return;
431      
432        yield* this.inOrder(rootNode.left!);
433        yield rootNode.data;
434        yield* this.inOrder(rootNode.right!);
435    }
436
437    // Generator for reverse in-order traversal
438    *reverseInOrder(rootNode: RBTreeNode<T> = this.root!): Generator<T, undefined, void> {
439        if (!rootNode) return;
440      
441        yield* this.reverseInOrder(rootNode.right!);
442        yield rootNode.data;
443        yield* this.reverseInOrder(rootNode.left!);
444    }
445}
446
447// TreeSet implementation using Red-Black Tree
448class TreeSet<T = number> {
449    private _size: number;
450    private tree: RBTree<T>;
451    private compare: CompareFunction<T>;
452
453    constructor(
454        collection: T[] | CompareFunction<T> = [],
455        compare: CompareFunction<T> = (left: T, right: T) => (left < right ? -1 : left > right ? 1 : 0)
456    ) {
457        // Handle overloaded constructor
458        if (typeof collection === 'function') {
459            compare = collection;
460            collection = [];
461        }
462      
463        this._size = 0;
464        this.compare = compare;
465        this.tree = new RBTree(compare);
466      
467        // Add initial collection
468        for (const value of collection) {
469            this.add(value);
470        }
471    }
472
473    // Get the size of the set
474    size(): number {
475        return this._size;
476    }
477
478    // Check if value exists in the set
479    has(value: T): boolean {
480        return !!this.tree.find(value);
481    }
482
483    // Add a value to the set
484    add(value: T): boolean {
485        const wasSuccessful = this.tree.insert(value);
486        if (wasSuccessful) {
487            this._size++;
488        }
489        return wasSuccessful;
490    }
491
492    // Delete a value from the set
493    delete(value: T): boolean {
494        const wasDeleted = this.tree.deleteAll(value);
495        if (wasDeleted) {
496            this._size--;
497        }
498        return wasDeleted;
499    }
500
501    // Find the smallest value >= given value
502    ceil(value: T): T | undefined {
503        let currentNode = this.tree.root;
504        let result = null;
505      
506        while (currentNode) {
507            if (this.compare(currentNode.data, value) >= 0) {
508                result = currentNode;
509                currentNode = currentNode.left;
510            } else {
511                currentNode = currentNode.right;
512            }
513        }
514      
515        return result?.data;
516    }
517
518    // Find the largest value <= given value
519    floor(value: T): T | undefined {
520        let currentNode = this.tree.root;
521        let result = null;
522      
523        while (currentNode) {
524            if (this.compare(value, currentNode.data) >= 0) {
525                result = currentNode;
526                currentNode = currentNode.right;
527            } else {
528                currentNode = currentNode.left;
529            }
530        }
531      
532        return result?.data;
533    }
534
535    // Find the smallest value > given value
536    higher(value: T): T | undefined {
537        let currentNode = this.tree.root;
538        let result = null;
539      
540        while (currentNode) {
541            if (this.compare(value, currentNode.data) < 0) {
542                result = currentNode;
543                currentNode = currentNode.left;
544            } else {
545                currentNode = currentNode.right;
546            }
547        }
548      
549        return result?.data;
550    }
551
552    // Find the largest value < given value
553    lower(value: T): T | undefined {
554        let currentNode = this.tree.root;
555        let result = null;
556      
557        while (currentNode) {
558            if (this.compare(currentNode.data, value) < 0) {
559                result = currentNode;
560                currentNode = currentNode.right;
561            } else {
562                currentNode = currentNode.left;
563            }
564        }
565      
566        return result?.data;
567    }
568
569    // Get the smallest value in the set
570    first(): T | undefined {
571        return this.tree.inOrder().next().value;
572    }
573
574    // Get the largest value in the set
575    last(): T | undefined {
576        return this.tree.reverseInOrder().next().value;
577    }
578
579    // Remove and return the smallest value
580    shift(): T | undefined {
581        const firstValue = this.first();
582        if (firstValue === undefined) return undefined;
583      
584        this.delete(firstValue);
585        return firstValue;
586    }
587
588    // Remove and return the largest value
589    pop(): T | undefined {
590        const lastValue = this.last();
591        if (lastValue === undefined) return undefined;
592      
593        this.delete(lastValue);
594        return lastValue;
595    }
596
597    // Iterator implementation
598    *[Symbol.iterator](): Generator<T, void, void> {
599        yield* this.values();
600    }
601
602    // Get all keys (same as values for a set)
603    *keys(): Generator<T, void, void> {
604        yield* this.values();
605    }
606
607    // Get all values in ascending order
608    *values(): Generator<T, undefined, void> {
609        yield* this.tree.inOrder();
610        return undefined;
611    }
612
613    // Get all values in descending order
614    *rvalues(): Generator<T, undefined, void> {
615        yield* this.tree.reverseInOrder();
616        return undefined;
617    }
618}
619
620// TreeMultiSet implementation (allows duplicates)
621class TreeMultiSet<T = number> {
622    private _size: number;
623    private tree: RBTree<T>;
624    private compare: CompareFunction<T>;
625
626    constructor(
627        collection: T[] | CompareFunction<T> = [],
628        compare: CompareFunction<T> = (left: T, right: T) => (left < right ? -1 : left > right ? 1 : 0)
629    ) {
630        // Handle overloaded constructor
631        if (typeof collection === 'function') {
632            compare = collection;
633            collection = [];
634        }
635      
636        this._size = 0;
637        this.compare = compare;
638        this.tree = new RBTree(compare);
639      
640        // Add initial collection
641        for (const value of collection) {
642            this.add(value);
643        }
644    }
645
646    // Get the total size of the multiset
647    size(): number {
648        return this._size;
649    }
650
651    // Check if value exists in the multiset
652    has(value: T): boolean {
653        return !!this.tree.find(value);
654    }
655
656    // Add a value to the multiset (allows duplicates)
657    add(value: T): boolean {
658        this.tree.insert(value);
659        this._size++;
660        return true;
661    }
662
663    // Delete a single occurrence of the value
664    delete(value: T): boolean {
665        const wasSuccessful = this.tree.delete(value);
666        if (wasSuccessful) {
667            this._size--;
668        }
669        return wasSuccessful;
670    }
671
672    // Count occurrences of a value
673    count(value: T): number {
674        const node = this.tree.find(value);
675        return node ? node.count : 0;
676    }
677
678    // Find the smallest value >= given value
679    ceil(value: T): T | undefined {
680        let currentNode = this.tree.root;
681        let result = null;
682      
683        while (currentNode) {
684            if (this.compare(currentNode.data, value) >= 0) {
685                result = currentNode;
686                currentNode = currentNode.left;
687            } else {
688                currentNode = currentNode.right;
689            }
690        }
691      
692        return result?.data;
693    }
694
695    // Find the largest value <= given value
696    floor(value: T): T | undefined {
697        let currentNode = this.tree.root;
698        let result = null;
699      
700        while (currentNode) {
701            if (this.compare(value, currentNode.data) >= 0) {
702                result = currentNode;
703                currentNode = currentNode.right;
704            } else {
705                currentNode = currentNode.left;
706            }
707        }
708      
709        return result?.data;
710    }
711
712    // Find the smallest value > given value
713    higher(value: T): T | undefined {
714        let currentNode = this.tree.root;
715        let result = null;
716      
717        while (currentNode) {
718            if (this.compare(value, currentNode.data) < 0) {
719                result = currentNode;
720                currentNode = currentNode.left;
721            } else {
722                currentNode = currentNode.right;
723            }
724        }
725      
726        return result?.data;
727    }
728
729    // Find the largest value < given value
730    lower(value: T): T | undefined {
731        let currentNode = this.tree.root;
732        let result = null;
733      
734        while (currentNode) {
735            if (this.compare(currentNode.data, value) < 0) {
736                result = currentNode;
737                currentNode = currentNode.right;
738            } else {
739                currentNode = currentNode.left;
740            }
741        }
742      
743        return result?.data;
744    }
745
746    // Get the smallest value in the multiset
747    first(): T | undefined {
748        return this.tree.inOrder().next().value;
749    }
750
751    // Get the largest value in the multiset
752    last(): T | undefined {
753        return this.tree.reverseInOrder().next().value;
754    }
755
756    // Remove and return the smallest value
757    shift(): T | undefined {
758        const firstValue = this.first();
759        if (firstValue === undefined) return undefined;
760      
761        this.delete(firstValue);
762        return firstValue;
763    }
764
765    // Remove and return the largest value
766    pop(): T | undefined {
767        const lastValue = this.last();
768        if (lastValue === undefined) return undefined;
769      
770        this.delete(lastValue);
771        return lastValue;
772    }
773
774    // Iterator implementation
775    *[Symbol.iterator](): Generator<T, void, void> {
776        yield* this.values();
777    }
778
779    // Get all keys (includes duplicates)
780    *keys(): Generator<T, void, void> {
781        yield* this.values();
782    }
783
784    // Get all values including duplicates in ascending order
785    *values(): Generator<T, undefined, void> {
786        for (const value of this.tree.inOrder()) {
787            let occurrences = this.count(value);
788            while (occurrences--) {
789                yield value;
790            }
791        }
792        return undefined;
793    }
794
795    // Get all values including duplicates in descending order
796    *rvalues(): Generator<T, undefined, void> {
797        for (const value of this.tree.reverseInOrder()) {
798            let occurrences = this.count(value);
799            while (occurrences--) {
800                yield value;
801            }
802        }
803        return undefined;
804    }
805}
806
807// Initialize the infinite set on module load
808initializeSmallestInfiniteSet();
809

Time and Space Complexity

Time Complexity:

__init__(): O(n) where n = 1000, as it creates a SortedSet with 1000 elements
popSmallest(): O(log n) for accessing the smallest element and removing it from the SortedSet
addBack(): O(log n) for adding an element back to the SortedSet while maintaining sorted order

Space Complexity: O(n) where n represents the maximum number of elements stored in the set. Initially, the set stores 1000 elements (from 1 to 1000), so the space complexity is O(1000) = O(n). Even though elements can be popped and added back, the maximum space used is bounded by the initial range of values, making the overall space complexity O(n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Dependency on External Library

The solution relies on sortedcontainers.SortedSet, which isn't part of Python's standard library. This can cause issues in coding interviews or platforms that don't support external libraries.

Solution: Implement using standard library components:

import heapq

class SmallestInfiniteSet:
    def __init__(self):
        self.next_num = 1  # Track the next consecutive number
        self.added_back = []  # Min heap for numbers added back
        self.added_set = set()  # Track what's in the heap to avoid duplicates
  
    def popSmallest(self) -> int:
        if self.added_back:
            smallest = heapq.heappop(self.added_back)
            self.added_set.remove(smallest)
            return smallest
        else:
            smallest = self.next_num
            self.next_num += 1
            return smallest
  
    def addBack(self, num: int) -> None:
        if num < self.next_num and num not in self.added_set:
            heapq.heappush(self.added_back, num)
            self.added_set.add(num)

2. Memory Inefficiency with Pre-population

Pre-populating 1000 elements wastes memory when only a few operations are performed. If the problem constraints change to allow larger numbers, this approach becomes impractical.

Solution: Use a lazy approach that tracks only the numbers that have been removed or added back, using a pointer for the next consecutive number in the infinite sequence.

3. Not Handling Edge Cases in addBack

A common mistake is not checking whether the number being added back:

Was actually removed before (is less than the current smallest consecutive number)
Is already in the set (duplicate addition)

Solution: Always validate:

def addBack(self, num: int) -> None:
    # Only add if it was previously popped and not already added back
    if num < self.next_num and num not in self.added_set:
        heapq.heappush(self.added_back, num)
        self.added_set.add(num)

4. Incorrect Assumption About Number Range

Assuming numbers will always be within 1-1000 without verifying problem constraints can lead to failures if larger numbers are added back.

Solution: Either clarify constraints or implement a dynamic approach that doesn't rely on a fixed range, using the combination of a counter and heap as shown above.

Discover Your Strengths and Weaknesses: Take Our 5-Minute Quiz to Tailor Your Study Plan:

Which of these properties could exist for a graph but not a tree?

Contain cycles

There are multiple ways to traverse

Connected

Has no cycle

I don't know