[백준/이분탐색] 2243번 사탕상자 (JAVA)

문제

---

풀이

이분탐색을 풀려다가 세그먼트 트리까지 공부해버리게 만든 문제다.

• 세그먼트 트리는 구간합을 트리 형태로 저장해서, 값이 자주 변하는 경우 빠르게 부분합을 계산할 수 있도록 만든 자료구조다.
• 배열의 합을 선형적으로 직접 계산할 경우 O(MN)이지만, 세그먼트 트리는 O(logN)의 시간복잡도로 원하는 값을 탐색을 할 수 있다.
• 세그먼트 트리는 Full binary tree에 가깝기 때문에 실제 트리 클래스를 구현하지 않고 배열에 저장하는 것이 효율적이다.
• 트리를 저장하는 배열의 크기는, 원본 데이터가 N개인 경우, 2^k > N인 k에 대해 2^k * 2이다. 하지만 실제로는 간단하게 N*4로 설정해도 된다.

본 문제에서는 tree를 세그먼트 트리로 활용한다.
tree의 인덱스 값은 사탕의 맛 순위이고, 저장된 값은 사탕의 개수이다.

update()

1. node, start, end는 노드에 대한 정보를 담는다.
   • node: 노드의 번호
   • start: 노드가 저장하는 부분합의 시작값 인덱스
   • end: 노드가 저장하는 부분합의 끝값 인덱스
2. target: 업데이트 할 사탕의 맛 순위
3. count: 업데이트 할 개수

binarySearch()

1. node, start, end는 노드에 대한 정보를 담는다.
2. target: 찾고 싶은 사탕 번호

• 재귀적으로 이분탐색을 진행한다.
• start >= end가 되면 사탕의 맛 순위를 찾은 것이므로, 사탕을 하나 꺼내고 순위를 리턴한다.
• target 값과 현재 node의 자식 노드 값을 비교해서 탐색할 방향을 결정한다.

코드

import java.io.*;

public class Main {

    private static final int MAX = 1000000;
    private static final int TYPE = 0;
    private static final int PICK = 1;
    private static final int UPDATE = 2;

    private static int[] tree = new int[(MAX + 1) * 4];

    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());
        for (int i = 0; i < n; i++) {
            String[] line = br.readLine().split(" ");
            if (Integer.parseInt(line[TYPE]) == PICK)
                pick(Integer.parseInt(line[1]), bw);
            else if (Integer.parseInt(line[TYPE]) == UPDATE)
                update(1, 1, MAX, Integer.parseInt(line[1]), Integer.parseInt(line[2]));
        }

        br.close();
        bw.close();
    }

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

    private static void pick(int target, BufferedWriter bw) throws IOException {
        bw.append(String.valueOf(binarySearch(1, 1, MAX, target))).append("\n");
    }

    private static int binarySearch(int node, int start, int end, int target) {
        if (start >= end) {
            update(1, 1, MAX, start, -1);
            return start;
        }

        int mid = (start + end) / 2;
        if (tree[node * 2] >= target)
            return binarySearch(node * 2, start, mid, target);
        else
            return binarySearch(node * 2 + 1, mid + 1, end, target - tree[node * 2]);
    }
}

---

참고

boj 2243 사탕상자
세그먼트 트리 (Segment Tree)