[Java] LeetCode 2336: Smallest Number in Infinite Set

[문제 바로 가기] - Smallest Number in Infinite Set

문제 해석

양의 정수 [1, 2, 3, 4, 5, ...]가 무한히 존재하는 집합이 있다고 가정할 때, SmallestInfiniteSet 클래스를 구현해라.

• SmallestInfiniteSet() : 클래스 생성자로, 모든 양의 정수가 포함된 무한한 집합으로 초기화
• int popSmallest() : 현재 집합에서 가장 작은 수를 꺼내서 제거하고, 그 수를 반환
• void addBack(int num) : 정수 num이 현재 집합에 없다면 추가하고, 이미 있다면 아무 일도 일어나지 않음

---

💡 아이디어

popSmallest()를 위해 PriorityQueue를, addBack()을 위해 Set을 사용해야 함은 알았지만 구현 과정에서 오류가 있었다.

처음 풀이

import java.util.*;

/**
 * 리트코드 2336번 Smallest Number in Infinite Set
 * - Set에 addBack한 num 넣기
 * - PriorityQueue에 Set에 없는 값 넣기
 */

class SmallestInfiniteSet {
    Set<Integer> set; // pop된 정수
    PriorityQueue<Integer> pq;

    // 모든 양의 정수가 포함된 무한한 집합
    public SmallestInfiniteSet() {
        set = new HashSet<>();
        pq = new PriorityQueue<>();

        pq.offer(1);
    }
    
    // 가장 작은 값을 꺼내서 제거하고 반환 : PriorityQueue
    public int popSmallest() {
        int small = pq.poll();
        set.add(small);

        int index = 1;
        while (true) {
            int next = small + index;
            if (!set.contains(next)) {
                pq.offer(next);
                index++;
                break;
            } else index++;
        }

        return small;
    }
    
    // 현재 집합에 num이 없으면 다시 추가 : Set
    // 이미 있다면 아무 일도 일어나지 않음
    public void addBack(int num) {
        if (set.contains(num)) {
            pq.offer(num);
            set.remove(num);
        }
    }
}

• SmallestInfiniteSet() : 당연히 무한한 집합을 모두 PriorityQueue에 넣으면 안된다고 생각하고, 먼저 PriorityQueue에 1만 넣었다.

• popSmallest() : PriorityQueue에서 poll()하여 가장 작은 값을 삭제한 뒤 해당 값을 Set에 넣었다. 이후 다음으로 작은 값을 PriorityQueue에 넣는데 Set에 있는 값은 삭제된 값이므로, Set에 없는 가장 작은 수를 넣어줬다.

• addBack() : Set에 값이 존재한다면 현재 집합에 num이 존재하지 않다는 뜻이므로 PriorityQueue에 추가하고 Set에서는 지운다.

Input
["SmallestInfiniteSet", "addBack", "popSmallest", "popSmallest", "popSmallest", "addBack", "popSmallest", "popSmallest", "popSmallest"][], [2], [], [], [], [1], [], [], []]

Output
[null, null, 1, 2, 3, null, 1, 4, 5]

내 풀이는 이 예제에서 [null, null, 1, 2, 3, null, 1, 4, 4]가 된다. 즉, 4가 2번 들어가게 돼서 중복 숫자가 PriorityQueue에 들어가게 되는것이다.

최종 풀이

import java.util.*;

class SmallestInfiniteSet {
    int current; // 아직 pop되지 않은 가장 작은 자연수
    PriorityQueue<Integer> pq;  // addBack으로 다시 추가된 수들
    Set<Integer> set; // pq 중복 방지용

    public SmallestInfiniteSet() {
        current = 1;
        pq = new PriorityQueue<>();
        set = new HashSet<>();
    }
    
    public int popSmallest() {
        if (!pq.isEmpty()) {
            int smallest = pq.poll();
            set.remove(smallest);
            return smallest;
        } return current++;
    }
    
    public void addBack(int num) {
        if (num < current && set.add(num)) {
            pq.offer(num);
        }
    }
}

• SmallestInfiniteSet() : current 1로 초기화

• popSmallest() : PriorityQueue가 비어있지 않다면 poll()한 후 Set에서 삭제 / 비어있다면 현재 current 값 반환하고 ++

• addBack() : num이 current보다 작고 Set에 없을 경우 PriorityQueue에 넣기