PS/Binary Search Tree

백준 21944번: 문제 추천 시스템 Version 2 (JAVA)

닻과매 2022. 5. 3. 20:24

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

 

21944번: 문제 추천 시스템 Version 2

recommend, recommend2, recommend3 명령이 주어질 때마다 문제 번호를 한 줄씩 출력한다. 주어지는 recommend, recommend2, recommend3 명령어의 총 개수는 최소 1개 이상이다.

www.acmicpc.net

 


 

풀이

명령에는 총 5가지: 추천 3가지 / 추가 / 삭제가 있다.

내부의 모든 위치에 대해 정렬된 자료 구조가 필요해 보이므로, TreeSet을 사용하도록 한다. Tree 내부에 Problem이라는 class를 넣는다. Problem class는 난이도, 문제 번호, 알고리즘 종류를 변수로 가지며, 난이도에 대한 내림차순 이후 문제 번호에 대한 내림차순으로 정렬되도록 순서를 정한다.

 

 

각 명령어에 대한 처리 방법

recommend2 x(첫 번째는 나중에 살펴본다.)

x값에 따라 tree의 맨 마지막 원소 혹은 맨 첫 번째 원소의 문제 번호를 출력하면 된다: tree.last(), tree.first()라는 method를 활용한다.

 

recommend3 x L

문제 번호가 0이고 난이도가 L인 문제 P를 만들어서, x값에 따라 inf{x | x > P}, sup{x | x < P}를 출력하면 된다: tree.higher(), tree.lower()라는 method를 활용한다.

 

recommend G x

x에 따라 알고리즘 분류가 G인 문제 중 가장 큰 문제, 가장 작은 문제를 출력하면 된다. 다만, tree는 G에 대해 전혀 정렬되어 있지 않으므로, Problem의 정렬 기준에 따라 만든 tree에서 해당 문제를 찾는 데에는 O(NlogN)의 시간이 든다. 쿼리가 M(최대 10,000)개이므로, MNlogN > 10억 --> 시간 초과의 가능성이 존재한다.

따라서 알고리즘 별로 subtree를 만들어, 해당 알고리즘에 대응하는 subtree에서 최댓값/최솟값을 찾는다.

 

add P L G

처음에 input을 받을때처럼 추가해주면 된다.

 

solved P

추가해준 만큼 삭제하면 된다. 다만 P만 주어지기에, 우리는 P 값을 통해 대응하는 L, G값을 얻어낼 수 있어야 한다 --> Map 자료구조를 활용해야 한다. 따라서, 초기 값을 입력받을 때나 add할 때에 Map에도 (문제, {난이도, 알고리즘 분류}) 순서쌍을 넣어줘야 한다.

 

 

코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.StringTokenizer;
import java.util.TreeSet;

public class Main {

    public static void main(String[] args) throws IOException {
        StringBuilder sb = new StringBuilder();
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int N = Integer.parseInt(br.readLine());
        TreeSet<Problem> tree = new TreeSet<>();
        List<TreeSet<Problem>> subtrees = new ArrayList<>();
        for (int i = 0; i <= 100; i++) {
            subtrees.add(new TreeSet<Problem>());
        }
        HashMap<Integer, int[]> hMap = new HashMap<>();


        for (int i = 0; i < N; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            int no = Integer.parseInt(st.nextToken());
            int level = Integer.parseInt(st.nextToken());
            int group = Integer.parseInt(st.nextToken());
            subtrees.get(group).add(new Problem(no, level, group));
            tree.add(new Problem(no, level, group));
            hMap.put(no, new int[] {level, group});
        }


        int M = Integer.parseInt(br.readLine());
        for (int i = 0; i < M; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            String command = st.nextToken();
            if (command.equals("recommend")) {
                int group = Integer.parseInt(st.nextToken());
                int x = Integer.parseInt(st.nextToken());
                if (x == 1) {
                    sb.append(subtrees.get(group).last().no+"\n");
                } else {
                    sb.append(subtrees.get(group).first().no+"\n");
                }
            } else if (command.equals("recommend2")) {
                int x = Integer.parseInt(st.nextToken());
                if (x == 1) {
                    sb.append(tree.last().no+"\n");
                } else {
                    sb.append(tree.first().no+"\n");
                }
            } else if (command.equals("recommend3")) {
                int x = Integer.parseInt(st.nextToken());
                int level = Integer.parseInt(st.nextToken());
                if (x == 1) {
                    if (tree.ceiling(new Problem(0, level, 0)) == null) sb.append("-1\n");
                    else sb.append(tree.ceiling(new Problem(0, level, 0)).no+"\n");
                } else {
                    if (tree.lower(new Problem(0, level, 0)) == null) sb.append("-1\n");
                    else sb.append(tree.lower(new Problem(0, level, 0)).no+"\n");
                }
            } else if (command.equals("add")) {
                int no = Integer.parseInt(st.nextToken());
                int level = Integer.parseInt(st.nextToken());
                int group = Integer.parseInt(st.nextToken());
                subtrees.get(group).add(new Problem(no, level, group));
                tree.add(new Problem(no, level, group));
                hMap.put(no, new int[] {level, group});
            } else {
                int no = Integer.parseInt(st.nextToken());
                if (!hMap.containsKey(no)) continue;
                int level = hMap.get(no)[0];
                int group = hMap.get(no)[1];
                hMap.remove(no);
                tree.remove(new Problem(no, level, group));
                subtrees.get(group).remove(new Problem(no, level, group));
            }
        }
        System.out.println(sb.toString());
    }


    static class Problem implements Comparable<Problem> {
        int no;
        int level;
        int group;

        public Problem(int no, int level, int group) {
            this.no = no;
            this.level = level;
            this.group = group;
        }

        @Override
        public int compareTo(Problem o) {
            if (level == o.level) {
                return Integer.compare(no, o.no);
            }
            return Integer.compare(level, o.level);
        }

        @Override
        public boolean equals(Object obj) {
            if (this == obj)
                return true;
            if (obj == null)
                return false;
            if (getClass() != obj.getClass())
                return false;
            Problem other = (Problem) obj;
            if (group != other.group)
                return false;
            if (level != other.level)
                return false;
            if (no != other.no)
                return false;
            return true;
        }
    }

}