LeetCode2336无限集中的最小数字

题目描述

  现有一个包含所有正整数的集合 [1, 2, 3, 4, 5, …] 。实现 SmallestInfiniteSet 类：SmallestInfiniteSet() 初始化 SmallestInfiniteSet 对象以包含 所有 正整数。int popSmallest() 移除 并返回该无限集中的最小整数。void addBack(int num) 如果正整数 num 不 存在于无限集中，则将一个 num 添加 到该无限集最后。

解析

  用一个数字表示前面最大的数字，然后用最小堆或者红黑树去记录后面的数字即可。

class SmallestInfiniteSet {

        private int frontMax;
        private TreeSet<Integer> backElements;

        public SmallestInfiniteSet() {
            this.backElements = new TreeSet<>();
            this.frontMax = 1;
        }

        public int popSmallest() {
            int res = 0;
            if(backElements.isEmpty()) {
                res = frontMax;
                frontMax ++;
            }
            else{
                int backMin = backElements.first();
                if(backMin < frontMax) {
                    res = backMin;
                    backElements.remove(backMin);
                }
                else {
                    res = frontMax;
                    frontMax ++;
                }
            }
            return res;
        }

        public void addBack(int num) {
            if(num >= frontMax) {
                return;
            }
            backElements.add(num);
        }
    }