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

문제

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

2243번: 사탕상자

 

설명

각각의 사탕은 그 맛의 좋고 나쁨이 1부터 1,000,000까지의 정수로 구분된다.
1이 가장 맛있는 사탕을 의미하며, 1,000,000은 가장 맛없는 사탕을 의미한다.

A가 1인 경우에는 순위가 B인 사탕을 꺼내는 것이고 2인 경우는 맛이 B인 사탕을 C만큼 넣거나 빼는 것이다.

이를 구현하기 위해 인덱스 트리를 사용한다.

if (A == 1) {
    long B = Integer.parseInt(st.nextToken());
    int taste = query(1, max, 1, B);	// 꺼낸 사탕의 종류
    System.out.println(taste);
    update(1, max, 1, taste, -1);	// 해당 사탕 하나 꺼냄
}
else {
    int B = Integer.parseInt(st.nextToken());
    int C = Integer.parseInt(st.nextToken());
    update(1, max, 1, B, C);
}

query 메서드에서는 순위가 target인 사탕이 어디에 위치해있는지 루트 노드에서부터 하나씩 내려가면서 찾을 수 있다. 계속 내려가다가 left == right 가 되는 순간에 해당 순위인 사탕 맛을 리턴한다.

int query(int left, int right, int node, long target) {
        // 리프 노드인 경우 (사탕 찾음)
        if (left == right) {
            return left;
        }

        // 내부 노드인 경우
        int mid = (left + right) / 2;
        if (tree[node * 2] >= target) { // 왼쪽 노드이면 target
            return query(left, mid, node * 2, target);
        }
        else {  // 오른쪽 노드이면 (target - 왼쪽노드)
            return query(mid + 1, right, node * 2 + 1, target - tree[node * 2]);
        }
    }

update 메서드에서는 해당 노드에 target이 포함되지 않으면 바로 리턴하여 무시하고 포함되면 현재 노드에 diff를 반영한다. 자식이 있을 경우 왼쪽 자식과 오른쪽 자식 모두 update를 호출하여 diff를 반영한다.

void update(int left, int right, int node, int target, int diff) {
    // 1. 연관 없음
    if (target < left || right < target) {
        return;
    }

    // 2. 연관 있음
    tree[node] += diff;

    if (left != right) {
        int mid = (left + right) / 2;
        update(left, mid, node * 2, target, diff);
        update(mid + 1, right, node * 2 + 1, target, diff);
    }
}

 

전체 코드

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

public class Main {

    static int N, S;
    static int max = 1000000;
    static long[] tree;

    static int query(int left, int right, int node, long target) {
        // 리프 노드인 경우 (사탕 찾음)
        if (left == right) {
            return left;
        }

        // 내부 노드인 경우
        int mid = (left + right) / 2;
        if (tree[node * 2] >= target) { // 왼쪽 노드이면 target
            return query(left, mid, node * 2, target);
        }
        else {  // 오른쪽 노드이면 (target - 왼쪽노드)
            return query(mid + 1, right, node * 2 + 1, target - tree[node * 2]);
        }
    }

    static void update(int left, int right, int node, int target, int diff) {
        // 1. 연관 없음
        if (target < left || right < target) {
            return;
        }

        // 2. 연관 있음
        tree[node] += diff;

        if (left != right) {
            int mid = (left + right) / 2;
            update(left, mid, node * 2, target, diff);
            update(mid + 1, right, node * 2 + 1, target, diff);
        }
    }

    public static void main(String[] args) throws IOException {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));

        StringTokenizer st = new StringTokenizer(br.readLine());
        N = Integer.parseInt(st.nextToken());

        S = 1;
        while (S < max) {
            S *= 2;
        }

        tree = new long[S * 2];

        for (int i = 0; i < N; i++) {
            st = new StringTokenizer(br.readLine());
            int A = Integer.parseInt(st.nextToken());

            if (A == 1) {
                long B = Integer.parseInt(st.nextToken());
                int taste = query(1, max, 1, B);    // 꺼낸 사탕의 종류
                System.out.println(taste);
                update(1, max, 1, taste, -1);   // 해당 사탕 하나 꺼냄
            }
            else {
                int B = Integer.parseInt(st.nextToken());
                int C = Integer.parseInt(st.nextToken());
                update(1, max, 1, B, C);
            }
        }
    }
}