[백준, 자바] 21939번 문제 추천 시스템 Version 1, TreeSet

1. 문제

https://www.acmicpc.net/problem/21939

 

2. 코드

• Problem의 정렬 조건 
  • 1. 난이도가 낮은 순으로 정렬한다.
  • 2. 난이도가 같다면 문제 번호가 낮은 순으로 정렬한다.
• TreeSet을 사용하여 Problem 저장 시에 정렬이 될 수 있도록 한다.
• TreeSet은 first(), last()를 사용하여 첫 번째 요소와 마지막 요소에 접근할 수 있다. (양방향 접근)

public class Main {

    static int N, M;
    static TreeSet<Problem> set = new TreeSet<>();
    static Map<Integer, Problem> map = new HashMap<>(); // {문제 번호, 문제}

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        N = Integer.parseInt(br.readLine());
        for (int i = 0; i < N; i++) {
            String[] inputs = br.readLine().split(" ");
            int problemNumber = Integer.parseInt(inputs[0]); // 문제 번호
            int problemDifficulty = Integer.parseInt(inputs[1]); // 문제 난이도
            Problem problem = new Problem(problemNumber, problemDifficulty);
            set.add(problem);
            map.put(problemNumber, problem);
        }

        StringBuilder sb = new StringBuilder();
        M = Integer.parseInt(br.readLine());
        for (int i = 0; i < M; i++) {
            String[] inputs = br.readLine().split(" ");
            if (inputs.length == 3) { // add
                int problemNumber = Integer.parseInt(inputs[1]);
                int problemDifficulty = Integer.parseInt(inputs[2]);
                Problem problem = new Problem(problemNumber, problemDifficulty);
                set.add(problem);
                map.put(problemNumber, problem);
            } else {
                if ("solved".equals(inputs[0])) { // solved
                    int problemNumber = Integer.parseInt(inputs[1]);
                    Problem removedProblem = map.remove(problemNumber);
                    set.remove(removedProblem);
                } else { // recommend
                    int x = Integer.parseInt(inputs[1]);
                    // TreeSet을 사용하면 양방향으로 접근 가능
                    if (x == 1) { // 가장 어려운 문제 번호 출력
                        Problem problem = set.last(); // TreeSet#last() : 마지막 요소 접근
                        sb.append(problem.number + "\n");
                    } else { // 가장 쉬운 문제 번호 출력
                        Problem problem = set.first(); // TreeSet#first() : 첫번째 요소 접근
                        sb.append(problem.number + "\n");
                    }
                }
            }
        }

        System.out.println(sb);
    }

    static class Problem implements Comparable<Problem> {
        int number; // 문제 번호
        int difficulty; // 난이도

        public Problem(int number, int difficulty) {
            this.number = number;
            this.difficulty = difficulty;
        }

        @Override
        public int compareTo(Problem problem) {
            if (this.difficulty == problem.difficulty) {
                return this.number - problem.number;
            }
            return this.difficulty - problem.difficulty;
        }
    }
}

 

3. 새로 알게된 점

• TreeSet을 사용하면 first()와 last()를 사용하여 양방향 접근이 가능하다는 것
  • first(), last() 시간 복잡도 : O(1)
  • TreeSet은 요소를 삽입할 때 트리를 자동으로 정렬하므로, 가장 작은 요소와 가장 큰 요소를 트리의 최좌단과 최우단에 미리 참조하고 있다.
  • 따라서 first()와 last()는 추가적인 탐색이나 연산 없이, 단순히 참조 값을 반환하므로 시간 복잡도는 O(1)이다.