[백준 - 2243] 사탕상자 - Java

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

2243번: 사탕상자

 

n번째 맛의 사탕을 꺼내거나, 넣는 경우이다. n번째 맛의 사탕의 개수는 변동된다.

사탕의 맛은 최대 1,000,000까지 있고, 사탕상자에는 최대 100,000번 손을 대고, 사탕의 최대 개수는 2,000,000,000이다.

 

인덱스 트리로 응용하면 아주 쉬운 문제이다. 

첫 번째로 맛있는 사탕이 2개, 세 번째로 맛있는 사탕이 3개 들어있을 때, 사탕상자에서 네 번째로 맛있는 사탕을 하나 꺼낸다면 세 번째로 맛있는 사탕을 하나 꺼내게 된다. 

현재 left = 1, right = 4라고 했을 때, 왼쪽 자식보다 target( = 4 )이 더 크기 때문에, 네 번째 맛 사탕은 오른쪽 자식으로 타고 들어가야 한다. 따라서 오른쪽 자식으로 query를 보낸다. 이 때, target 값은 target-leftvalue로 갱신한다. 

 

 

제출하기 전 테스트 할 때는 S를 꼭 1,000,000와 같거나 큰 2의 제곱수로 설정하지 않고, 적당히 작은 값으로 설정해 테스트 하는 것도 좋은 방법이다.

 

public class Main {
    static long[] tree; // 예제는 크기를 줄여서 해보고 제출할 때 크게 해서 제출
    public static void main(String[] args) throws Exception {
        BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
        int N = Integer.parseInt(br.readLine());
        int S = 1;
        while(S<1000000) S*=2;
        tree = new long[2*S+1];

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

            if(A==1){ // B번째 순위의 사탕을 꺼내는 경우
                // B맛의 사탕의 index를 구함
                int index = query(1, S, 1, B)-S+1;
                System.out.println(index);
                // B맛의 사탕을 하나 꺼냄
                update(1, S, 1, index, -1);
            }
            if(A==2){ // B정도의 맛의 사탕을 넣거나 빼는 경우
                int C = Integer.parseInt(st.nextToken());
                // B맛의 사탕을 C개 넣음
                update(1, S, 1, B, C);

            }


        }

    }

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

        }
    }

    static int query(int left, int right, int node, long candy){
        // leafnode라면 node값과 비교
        if(left==right && tree[node]>=candy){
            return node;
        }
        else {
            int mid = (left+right)/2;
            long leftValue = tree[2*node];
            // 왼쪽에 있음
            if(leftValue >= candy){
                return query(left, mid, 2*node, candy);
            }
            // 오른쪽에 있음
            else {
                return query(mid+1, right, 2*node+1, candy-leftValue);
            }
        }
    }
}