[백준 2243] 사탕상자 / 자바 / 자료구조(세그먼트 트리)

#문제         

레벨: P5
알고리즘: 세그먼트 트리
풀이시간: 1시간
힌트 참조 유무: 유

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

---

#문제 풀이        

 

위 문제는 2가지로 구현할 것이다.

위 문제를 봤을 때 가장 직관적으로 떠오르는 자료구조는 TreeMap이다. 입력하면 항상 키의 오름차순대로 정렬된다. 사탕을 넣고 뺐을 때는 O(logM)으로 사탕을 줄 때는 O(M)이다. 왜 이렇게 되냐 하면, 키의 위치를 찾는 건 O(logM)이지만, 사탕의 순위를 알기 위해서는 1순위부터 차례대로 순위를 매겨가며 해당 순위의 사탕을 줘야 하기 때문에,  1~M순위까지의 Key를 다 탐색해야 하기 때문이다. 

 

그러나 해당문제는 세그먼트 트리 문제이다. 글보다는 그림이 이해하기 쉬우니 아래 그림을 봐보기 바란다. 세그먼트 트리는 하위 노드의 합이 들어있다.

초기 상태: 빈 트리

 

       [0]
    /       \
  [0]       [0]
 /   \     /   \
[0] [0] [0]    [0]

2 1 2 (1번 맛 사탕 2개 추가)
 

       [2]
    /       \
  [2]       [0]
 /   \     /   \
[2] [0] [0]   [0]

2 3 3 (3번 맛 사탕 3개 추가)
 

       [5]
    /       \
  [2]       [3]
 /   \     /   \
[2] [0] [3]    [0]

1 2 (2번째로 맛있는 사탕 꺼내기)
트리를 탐색하여 2번째 사탕이 1번 맛임을 찾음
1번 맛 사탕을 하나 제거
 

       [4]
    /       \
  [1]       [3]
 /   \     /   \
[1] [0] [3]    [0]

1 2 (2번째로 맛있는 사탕 꺼내기)
트리를 탐색하여 2번째 사탕이 3번 맛임을 찾음
3번 맛 사탕을 하나 제거
 

       [3]
    /       \
  [1]       [2]
 /   \     /   \
[1] [0] [2]    [0]

2 1 -1 (1번 맛 사탕 1개 제거)
 

       [2]
    /       \
  [0]       [2]
 /   \     /   \
[0] [0] [2]    [0]

1 2 (2번째로 맛있는 사탕 꺼내기)
트리를 탐색하여 2번째 사탕이 3번 맛임을 찾음
3번 맛 사탕을 하나 제거
 

       [1]
    /       \
  [0]       [1]
 /   \     /   \
[0] [0] [1]    [0]

 

참고로 세그먼트트리 코드가 TreeMap코드보다 10배는 빠르다. 세그먼트트리는 시간복잡도가 O(logM)이지만, TreeMap의 코드는 O(logM) + O(M)이기 때문이다.

---

#풀이 코드      

세그먼트 트리 코드

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

public class Main {
    static int[] tree;
    static int MAX = 1000000;

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        
        int n = Integer.parseInt(br.readLine());
        int size = 1;
        while (size < MAX) size *= 2;
        tree = new int[size * 2];

        for (int i = 0; i < n; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            int A = Integer.parseInt(st.nextToken());
            if (A == 1) {// 사탕을 줄 때
                int B = Integer.parseInt(st.nextToken());
                int result = query(1, 1, MAX, B);
                bw.write(result + "\n");
                update(1, 1, MAX, result, -1);
            } else {    //사탕을 넣고 뺄 때 
                int B = Integer.parseInt(st.nextToken());
                int C = Integer.parseInt(st.nextToken());
                update(1, 1, MAX, B, C);
            }
        }
        
        bw.flush();
        bw.close();
        br.close();
    }

    static int query(int node, int start, int end, int k) {  // start와 end는 사탕의 맛 범위이다. 총 사탕의 개수의 범위가 아니니 헷갈리지 마길 바란다.
        if (start == end) return start;
        int mid = (start + end) / 2;
        if (k <= tree[node*2]) return query(node*2, start, mid, k);
        else return query(node*2+1, mid+1, end, k - tree[node*2]);
    }

    static void update(int node, int start, int end, int index, int diff) {
        if (index < start || index > end) return;
        tree[node] += diff;
        if (start != end) {
            int mid = (start + end) / 2;
            update(node*2, start, mid, index, diff);
            update(node*2+1, mid+1, end, index, diff);
        }
    }
}

 

TreeMap 코드

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

public class Main {
    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        
        int n = Integer.parseInt(br.readLine());
        TreeMap<Integer, Integer> candyMap = new TreeMap<>();
        
        for (int i = 0; i < n; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            int A = Integer.parseInt(st.nextToken());
            
            if (A == 1) {
                int B = Integer.parseInt(st.nextToken());
                int count = 0;
                int taste = 0;
                for (Map.Entry<Integer, Integer> entry : candyMap.entrySet()) {
                    count += entry.getValue();
                    if (count >= B) {
                        taste = entry.getKey();
                        break;
                    }
                }
                bw.write(taste + "\n");
                candyMap.put(taste, candyMap.get(taste) - 1);
                if (candyMap.get(taste) == 0) {
                    candyMap.remove(taste);
                }
            } else {
                int B = Integer.parseInt(st.nextToken());
                int C = Integer.parseInt(st.nextToken());
                candyMap.put(B, candyMap.getOrDefault(B, 0) + C);
                if (candyMap.get(B) == 0) {
                    candyMap.remove(B);
                }
            }
        }
        
        bw.flush();
        bw.close();
        br.close();
    }
}