[Baekjoon] 사탕상자

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

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

public class Main {
    static final int MAX_VALUE = 1_000_000;

    static StringBuilder answer = new StringBuilder();
    static int[] segmentTree = new int[MAX_VALUE * 4];// 사탕의 순위

    public static void main(String[] args) throws IOException {
        init();
        // sol();
        print();
    }

    static void print() throws IOException {
        BufferedWriter bw = new BufferedWriter(new OutputStreamWriter(System.out));
        bw.write(String.valueOf(answer));
        bw.close();
    }

    static void init() throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        StringTokenizer st;

        int n = Integer.parseInt(br.readLine());

        for (int i = 0; i < n; i++) {
            st = new StringTokenizer(br.readLine());

            int A = Integer.parseInt(st.nextToken());
            int B = Integer.parseInt(st.nextToken());
            if (A == 1) {
                // 사탕 꺼내기
                int rank = B;
                int flavor = getSegment(1, 1, MAX_VALUE, rank);
                updateSegment(1, 1, MAX_VALUE, flavor, -1);
                answer.append(flavor).append("\n");
                continue;
            }
            // 사탕 넣기
            int flavor = B;
            int count = Integer.parseInt(st.nextToken());

            updateSegment(1, 1, MAX_VALUE, flavor, count);
        }
    }

    static int getSegment(int node, int left, int right, int index) {
        if (left == right) {
            return left;
        }

        int mid = (left + right) / 2;

        if (index <= segmentTree[node * 2]) {
            return getSegment(node * 2, left, mid, index);
        } else {
            return getSegment(node * 2 + 1, mid + 1, right, index - segmentTree[node * 2]);
        }
    }

    static void updateSegment(int node, int left, int right, int index, int value) {
        if (index < left || index > right) {
            return;
        }

        if (left == right) {
            segmentTree[node] += value;
            return;
        }

        int mid = (left + right) / 2;

        updateSegment(node * 2, left, mid, index, value);
        updateSegment(node * 2 + 1, mid + 1, right, index, value);

        segmentTree[node] = segmentTree[node * 2] + segmentTree[node * 2 + 1];
    }
}

주요 개념

• 세그먼트 트리

풀이 방법

• A의 입력값에 따라 B맛 사탕 개수를 C만큼 더하거나updateSegment
  B 순위 사탕의 맛을 구하고 한개를 뺀다getSegment + updateSegment.

• int[] segmentTree 에서
  index -> 사탕의 맛
  segmentTree[index] -> 현재 트리 레벨까지의 모든 사탕의 개수

• 기본적인 세그먼트 트리 문제에서 사용되는 updateSegment를 그대로 활용한다.

• 순위를 구하기 위해서 항상 left 노드를 사용할 지, right 노드를 사용할 지 매번 선택한다.

  • 찾으려는 순위가 6일 경우

  • left 노드segmentTree[node*2]는 4

  • right 노드segmentTree[node*2+1]는 6일 경우, 조건식에 의해 오른쪽 노드를 선택한다.

  • 이를 리프 노드까지if(left==right) 탐색할 경우, 순위 6의 맛을 구할 수 있다.