[Algorithm] boj_21944 - TreeSet 활용

문제

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

접근 방법

문제에서 요구하는 조건은 다음과 같다.

1. 알고리즘별로 문제를 분류할수 있어야한다.

2. 난이도(1순위), idx(2순위) 기준으로 정렬하고, 양쪽으로 조회 가능해야한다.
   -> treeSet의 first(),last(),floor(),ceiling()

3. 요소를 추가 할 수 있어야한다.

4. idx로 요소를 삭제 할수 있어야한다.

정렬/분류 기준이 너무 많고, 삽입/삭제는 물론이고 특정 값 기준으로 중간 요소를 탐색할수도 있어야 하기 때문에 어떤 자료구조를 사용해야할지 어려웠다.

TreeSet 사용

treeSet에는 ceiling(), floor() 메서드가 존재한다. recommend3을 구현하기 위해서는 이 메서드들을 사용해야한다.

• ceiling(): 특정 값보다 크거나 같은 값 중, 가장 작은 값을 출력한다.
• floor(): 특정 값보다 작거나 같은 값 중, 가장 큰 값을 출력한다.

❗️higher()/lower() 을 사용하면 안된다

-> higher()/lower()는 특정 값을 포함하지 않고 크거나 작은 값을 출력한다.

자료구조 활용

각각 다음의 자료구조 사용이 필요하다.

• 전체 문제 리스트(난이도/문제번호)
• 알고리즘별 문제 리스트
• 문제 번호로 조회할수있는 자료구조

코드

import java.io.*;
import java.util.*;

public class Main {

    static int N, M;
    static TreeSet<Question> questions = new TreeSet<>();
    static Map<Integer, TreeSet<Question>> algorithms = new HashMap<>();
    static Map<Integer, Question> info = new HashMap<>();
    static StringBuilder sb = new StringBuilder();

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        N = Integer.parseInt(br.readLine());

        //input
        for (int i = 0; i < N; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            int idx = Integer.parseInt(st.nextToken());
            int level = Integer.parseInt(st.nextToken());
            int al = Integer.parseInt(st.nextToken());
            Question question = new Question(al, idx, level);

            questions.add(question);
            algorithms.computeIfAbsent(al, key -> new TreeSet<>())
                    .add(question);
            info.put(idx, question);
        }

        M = Integer.parseInt(br.readLine());

        //run command
        for (int i = 0; i < M; i++) {
            doCommand(br.readLine());
        }

        System.out.println(sb);
    }

    private static class Question implements Comparable<Question>{
        int algorithm;
        int idx;
        int level;

        Question(int algorithm, int idx, int level) {
            this.algorithm = algorithm;
            this.idx = idx;
            this.level = level;
        }

        @Override
        public int compareTo(Question q) {
            if (this.level == q.level) {
                return this.idx - q.idx;
            }

            return this.level - q.level;
        }
    }

    static void doCommand(String cmd) {
        StringTokenizer st = new StringTokenizer(cmd);
        String command = st.nextToken();

        switch (command) {
            case "recommend" -> handleRecommend(st);
            case "recommend2" -> handleRecommend2(st);
            case "recommend3" -> handleRecommend3(st);
            case "add" -> handleAdd(st);
            case "solved" -> handleSolved(st);
        }
    }

    private static void handleSolved(StringTokenizer st) {
        int idx = Integer.parseInt(st.nextToken());
        Question targetQuestion = info.get(idx);

        //remove
        algorithms.get(targetQuestion.algorithm).remove(targetQuestion);
        questions.remove(targetQuestion);
        info.remove(idx);
    }

    private static void handleAdd(StringTokenizer st) {
        int idx = Integer.parseInt(st.nextToken());
        int level = Integer.parseInt(st.nextToken());
        int al = Integer.parseInt(st.nextToken());
        Question question = new Question(al, idx, level);

        questions.add(question);
        algorithms.computeIfAbsent(al, key -> new TreeSet<>())
                .add(question);
        info.putIfAbsent(idx, question);
    }

    private static void handleRecommend3(StringTokenizer st) {
        int x = Integer.parseInt(st.nextToken());
        int L = Integer.parseInt(st.nextToken());
        Question std = new Question(-1, -1, L);

        if (x == 1) {
            Question ceiling = questions.ceiling(std);
            if (Objects.isNull(ceiling)) {
                sb.append(-1).append("\n");
            } else {
                sb.append(ceiling.idx).append("\n");
            }
        } else {
            Question floor = questions.floor(std);
            if (Objects.isNull(floor)) {
                sb.append(-1).append("\n");
            } else {
                sb.append(floor.idx).append("\n");
            }
        }
    }

    private static void handleRecommend2(StringTokenizer st) {
        int x = Integer.parseInt(st.nextToken());

        if (x == 1) {
            int idx = questions.last().idx;
            sb.append(idx).append("\n");
        } else {
            int idx = questions.first().idx;
            sb.append(idx).append("\n");
        }
    }

    private static void handleRecommend(StringTokenizer st) {
        int G = Integer.parseInt(st.nextToken());
        int x = Integer.parseInt(st.nextToken());
        TreeSet<Question> q = algorithms.get(G);

        if (x == 1) {
            int idx = q.last().idx;
            sb.append(idx).append("\n");
        } else {
            int idx = q.first().idx;
            sb.append(idx).append("\n");
        }
    }
}