[백준 / 골드4] 21939 문제 추천 시스템 Version 1 (Java)

TreeSet은 이진탐색트리 자료구조를 사용하는 컬렉션이다.

문제 보기

사용한 것

• 검색과 삭제 모두 O(logN)으로 수행하기 위한 이진탐색트리

풀이 방법

• TreeSet을 사용하기 위해 Problem의 우선순위를 정의한 compareTo()를 람다식으로 넘긴다.
• 커맨드에 따라 다음을 수행한다.
  • "recommend" : treeSet에서 first 혹은 last를 출력
  • "add" : treeSet에 원소 추가
  • "solved" : treeSet에서 원소 삭제

코드

public class Main {

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;
        TreeSet<Problem> treeSet = new TreeSet<>((o1, o2) -> {
            if (o1.l == o2.l) {
                return o2.p - o1.p;
            }
            return o2.l - o1.l;
        });
        Map<Integer, Integer> map = new HashMap<>();
        int n = Integer.parseInt(br.readLine());
        int p;
        int l;
        for (int i = 0; i < n; i++) {
            st = new StringTokenizer(br.readLine());
            p = Integer.parseInt(st.nextToken());
            l = Integer.parseInt(st.nextToken());
            treeSet.add(new Problem(p, l));
            map.put(p, l);
        }
        StringBuilder sb = new StringBuilder();
        int m = Integer.parseInt(br.readLine());
        String command;
        for (int i = 0; i < m; i++) {
            st = new StringTokenizer(br.readLine());
            command = st.nextToken();
            if ("recommend".equals(command)) {
                command = st.nextToken();
                if ("1".equals(command)) {
                    sb.append(treeSet.first().p).append(lineSeparator());
                } else {
                    sb.append(treeSet.last().p).append(lineSeparator());
                }
            } else if ("add".equals(command)) {
                p = Integer.parseInt(st.nextToken());
                l = Integer.parseInt(st.nextToken());
                treeSet.add(new Problem(p, l));
                map.put(p, l);
            } else {
                p = Integer.parseInt(st.nextToken());
                treeSet.remove(new Problem(p, map.get(p)));
                map.remove(p);
            }
        }
        System.out.println(sb);
    }
}

class Problem {

    public int p;
    public int l;

    public Problem(int p, int l) {
        this.p = p;
        this.l = l;
    }
}