[백준 2243] 사탕상자 (JAVA)

문제

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

풀이

세그먼트 트리를 이용하여 해결할 수 있습니다. 사탕의 맛이 1000000이 최대이므로 그에 맞춰 세그먼트 트리를 만들어줍니다. 순위를 찾는 경우 현재 순위를 보고 자식 노드들의 값을 보고 판단합니다.

• 찾고 싶은 순위가 12위 인 경우 먼저 현재 노드의 왼쪽 자식 노드를 확인하여 왼쪽 노드의 값보다 작거나 같으므로 왼쪽 노드로 이동합니다.

• 현재 노드의 왼쪽 자식 노드보다 값이 크므로 지금 순위값에서 왼쪽 자식 노드를 빼서 오른쪽 자식으로 이동합니다.

• 리프 노드이므로 3번 맛을 출력합니다.

• 18번째 맛을 찾기 위해선 위에 사진처럼 진행됩니다.

코드

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

public class Main {

    public static void main(String[] args) throws IOException {

        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringBuilder sb = new StringBuilder();

        CandyBox candyBox = new CandyBox();
        int N = Integer.parseInt(br.readLine());
        for (int i = 0; i < N; i++) {
            StringTokenizer st = new StringTokenizer(br.readLine());
            if (Integer.parseInt(st.nextToken()) == 1) {
                sb.append(candyBox.getCandy(Integer.parseInt(st.nextToken()))).append('\n');
            } else {
                candyBox.putCandy(Integer.parseInt(st.nextToken()), Integer.parseInt(st.nextToken()));
            }
        }

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

    public static class CandyBox {
        int size;
        int[] candy, tree;

        public CandyBox() {
            int len = 1 << 20;
            this.candy = new int[1000001];
            this.tree = new int[len * 2];
            this.size = 1000000;
        }

        public void putCandy(int taste, int num) {
            update(1, size, 1, taste, num);
            candy[taste] += num;
        }

        public int getCandy(int rank) {
            int result = update(1, size, 1, rank);
            candy[result]--;
            return result;
        }

        public void update(int start, int end, int node, int taste, int num) {
            if (taste < start || end < taste) return;
            tree[node] += num;
            if (start == end) return;

            int mid = (start + end) >> 1;
            update(start, mid, node * 2, taste, num);
            update(mid + 1, end, node * 2 + 1, taste, num);
        }

        public int update(int start, int end, int node, int rank) {
            tree[node]--;
            if (start == end) return start;

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