[백준/Java] 2243 사탕상자

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

문제 요약

사탕의 맛은 1부터 100만의 숫자로 나타나며, 작을수록 맛이 좋다.
사탕을 꺼내는 경우에는 사탕의 맛 순위를 기반으로 꺼낸다.
사탕을 넣을 때는 그 맛과 개수를 알려준다. 사탕을 빼는 경우도 있다.
사탕을 꺼낼 때, 그 맛의 번호를 모두 출력하자.
손을 대는 횟수 n은 최대 10만이다. 사탕의 총 개수는 20억을 넘지 않는다.

---

문제 접근

이 문제도 인덱스 트리를 이용해서 해결할 수 있다.
인덱스 트리의 크기는 100만보다 큰, 가장 작은 2의 제곱수의 두 배로 한다.
합의 인덱스 트리로 구성하며, 왼쪽부터 맛있는 사탕이 있다.
사탕이 들어오면 해당 리프에 그 개수를 넣는다. 빼는 경우도 리프에서 빼면 된다.
사탕을 꺼내는 경우에는 루트에서 시작해서, 그 자식을 봐야 한다.
왼쪽 자식(왼쪽에 있는 개수)가 그 순위보다 크다면 왼쪽으로 가고, 아니면 오른쪽으로 간다.
오른쪽으로 가는 경우에는 그 순위를 왼쪽 자식의 값으로 뺀 다음, 가야 한다.
이를 반복하면 해당하는 사탕에 도달할 수 있다.

---

풀이

기본적인 아이디어는 다음과 같다.

1. (정리 예정)

이를 구현한 코드는 다음과 같다.

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

public class Main {
	
	static int DELIGHT_RANGE = 1000000;
	
	static int n, leafPointer;
	
	
	static int[] delightIndexTree;
	
	// 인덱스 트리 구성하기
	public static void makeTree() {
		
		// 리프 포인터, 트리 크기 구하기
		leafPointer = 1;
		
		while (leafPointer < DELIGHT_RANGE) {
			leafPointer <<= 1; 
		}
		
		int treeSize = leafPointer << 1;
		
		// 맛 인덱스 트리 초기화
		delightIndexTree = new int[treeSize];
		
		for (int i = 0; i < treeSize; i++) {
			delightIndexTree[i] = 0;
		}
	}
    
    // 사탕을 넣거나 꺼내는 경우의 함수
	// 여기서도 맛이 1부터 시작한다는 점 참고
	public static void updateCandy(int delight, int	quantity) {
		
		int targetIndex = delight + leafPointer - 1;
		
		// 리프 수정
		delightIndexTree[targetIndex] += quantity;
		targetIndex >>>= 1;
		
		// 리프의 부모를 따라가면서 수정
		while (targetIndex > 0) {
			delightIndexTree[targetIndex] = delightIndexTree[targetIndex*2] + delightIndexTree[targetIndex*2+1];
			targetIndex >>>= 1;
		}
	}
    
    
    // 사탕을 빼는 경우의 함수
	public static int takeCandy(int rank) {

		int nowIndex = 1;
		int delight;
		
		// 트리를 벗어날 때까지, 왼쪽 자식의 값을 rank와 비교하며 리프를 찾는다.
		while (true) {
			
			int leftChild = nowIndex * 2;
			int rightChild = nowIndex * 2 + 1;
			
			if (leftChild >= 2 * leafPointer) {
				break;
			}
			
			if (rank <= delightIndexTree[leftChild]) {
				nowIndex = leftChild;
			}
			// 왼쪽 자식보다 큰 값이라면 오른쪽 자식에 있다.
			// 이때 순위를 왼쪽 자식의 값을 빼서 보정해야 한다.
			else {
				nowIndex = rightChild;
				rank -= delightIndexTree[leftChild];
			}
		}
		
		delight = nowIndex - leafPointer + 1;
		
		// 하나 꺼냈으니까 업데이트 해야된다.
		updateCandy(delight, -1);
		
		return delight;
	}
    
    
    public static void main(String[] args) throws IOException {
    	
    	BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
    	StringBuilder sb = new StringBuilder();
    	
        // n 받기
    	n = Integer.parseInt(br.readLine());
        
        // 인덱스 트리 구성
        makeTree();
        
        
        // 명령을 받으면서 문제 해결하기
        // 사탕을 꺼낼 때는 그 맛을 출력도 해야 한다(순위가 아니라 맛을 나타내는 정수).
        for (int i = 0; i < n; i++) {
        	StringTokenizer st = new StringTokenizer(br.readLine());
        	
        	int order = Integer.parseInt(st.nextToken());
        	// 사탕을 꺼내서 주는 경우
        	if (order == 1) {
        		int rank = Integer.parseInt(st.nextToken());
        		
        		int delight = takeCandy(rank);
        		sb.append(delight+"\n");
        		//System.out.println(delight);
        	}
        	// 사탕을 넣거나 빼는 경우
        	else if (order == 2) {
        		int delight = Integer.parseInt(st.nextToken());
        		int quantity = Integer.parseInt(st.nextToken());
        		updateCandy(delight, quantity);
        	}
        }
        
        System.out.println(sb);
        
    }
}

---

회고

틀렸던 이유
인덱스 트리 테스트 용으로 크기를 작게 설정해놓고 복원하지 않아 틀렸다.

틀린 것은 아니지만, 값을 System.out.println으로 하나하나 출력하는 것보다 StringBuilder에 기록하고 마지막에 한 번에 출력하는 것이 시간적으로 더 효율적이라는 점을 확인했다.