백준 21939 - 문제 추천 시스템 Version 1

문제

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

리뷰

자바에서 키를 기준으로 키-값 쌍 데이터를 정렬된 형태로 저장하는 TreeMap을
이용하여 풀이할 수 있는 문제였다.

주어진 문제에서 구현함에 있어 가장 까다로운 연산이 solved 였던 것 같다.
단순히 정렬 기준이 되는 난이도를 기준으로 PriorityQueue등에 문제를
저장하였을 때, 해당 데이터를 삭제하려면 $O(N)$의 시간이 소요되기 때문에
무조건 시간 초과가 발생할 수 밖에 없었다. 따라서 삭제 연산의 복잡도가
$O(logN)$이하가 되도록 구현해야 했다.

로직을 구성하기 위해 우선, 문제를 나타내는 클래스인 Problem을 정의하였다.
그리고 다음 두 개의 TreeMap을 이용하였다.

• TreeMap<Problem, Integer> map
Problem 객체를 키로 하여 문제 번호를 값으로 저장한다. Comparator를
  통해 문제 난이도가 같을 때는 더 번호가 작은 것이 앞으로 정렬되게 하여
  firstKey()를 통해 접근했을 시 난이도와 번호가 가장 작은 데이터가,
  lastKey()를 통해 접근했을 시 난이도와 번호가 가장 높은 데이터가
  탐색될 수 있도록 구성했다.
• TreeMap<Integer, Problem> map2
  문제 번호를 키로 하여 Problem 인스턴스를 값으로 저장한다.

위 두 TreeMap을 이용하여 연산은 다음과 같은 논리로 구현하였다.

• recommend x
map의 firstKey()를 통해 난이도가 가장 낮은 문제를, lastKey() 연산을
  통해 난이도가 가장 높은 문제를 추천한다. Comparator로 번호 조건도 충족하였다.

• add P L
Problem 인스턴스를 생성한 후, map와 map2에 저장한다.

• solved P
  문제 번호가 P인 인스턴스를 탐색하기 위해 map2 를 이용한다. 이후 해당
  인스턴스를 map에서 조회하여 삭제한다. TreeMap에서 삭제 연산은 평균적으로
  $O(logN)$ 안에 이뤄지므로, 시간 제한 조건을 충족시킬 수 있다.

유의할 점은 TreeMap 내에서 Problem을 비교하기 때문에 equals 메서드를
적절히 오버라이딩해줄 필요가 있었다는 점이다.

로직의 시간복잡도는 명령을 처리하는 부분에서 $O(MlogN)$으로 수렴하므로,
$N=100,000$, $M=10,000$인 최악의 경우에도 무난히 제한 조건 1초를 통과한다.

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;
import java.util.TreeMap;

import static java.lang.Integer.compare;
import static java.lang.Integer.parseInt;

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        TreeMap<Problem, Integer> map = new TreeMap<>((p1, p2) -> {
            if (p1.score == p2.score) return compare(p1.num, p2.num);

            return compare(p1.score, p2.score);
        });
        TreeMap<Integer, Problem> map2 = new TreeMap<>();

        int N = parseInt(br.readLine());

        Problem p;
        int score, num;
        StringTokenizer st;
        while (N-- > 0) {
            st = new StringTokenizer(br.readLine());
            num = parseInt(st.nextToken());
            score = parseInt(st.nextToken());

            p = new Problem(score, num);
            map.put(p, num);
            map2.put(num, p);
        }


        int M = parseInt(br.readLine());
        Problem key;
        String cmd;
        int x;
        StringBuilder sb = new StringBuilder();
        while (M-- > 0) {
            st = new StringTokenizer(br.readLine());
            cmd = st.nextToken();

            switch (cmd) {
                case "add":
                    num = parseInt(st.nextToken());
                    score = parseInt(st.nextToken());
                    p = new Problem(score, num);
                    map.put(p, num);
                    map2.put(num, p);
                    break;
                case "solved":
                    num = parseInt(st.nextToken());
                    p = map2.get(num);
                    map.remove(p);
                    break;
                case "recommend":
                    x = parseInt(st.nextToken());
                    key = x == -1 ? map.firstKey() : map.lastKey();
                    sb.append(map.get(key)).append("\n");
                    break;
            }
        }

        System.out.print(sb);
        br.close();
    }

    static class Problem {
        int score, num;

        public Problem(int score, int num) {
            this.score = score;
            this.num = num;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            Problem problem = (Problem) o;
            return score == problem.score && num == problem.num;
        }
    }
}

결과