2336. Smallest Number in Infinite Set 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
class SmallestInfiniteSet {

    private final TreeSet<Integer> treeSet; // 存储不存在于无限集的数字

    public SmallestInfiniteSet() {
        this.treeSet = new TreeSet<>();
    }

    public int popSmallest() {
        // 寻找第一个缺失的数字

        int prev = 0;
        Integer higher = treeSet.higher(prev);

        while (higher != null && higher == prev + 1) {
            prev = higher;
            higher = treeSet.higher(prev);
        }

        // higher == null or higher != prev + 1
        treeSet.add(prev + 1);
        return prev + 1;
    }

    public void addBack(int num) {
        treeSet.remove(num);
    }
}

References

2336. Smallest Number in Infinite Set