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

#2243 사탕상자

난이도 : 플레 5

유형 : 세그먼트 트리 / 이분 탐색

2243번: 사탕상자

▸ 문제

수정이는 어린 동생을 달래기 위해서 사탕을 사용한다. 수정이는 평소에 여러 개의 사탕을 사서 사탕상자에 넣어두고, 동생이 말을 잘 들을 때면 그 안에서 사탕을 꺼내서 주곤 한다.

각각의 사탕은 그 맛의 좋고 나쁨이 1부터 1,000,000까지의 정수로 구분된다. 1이 가장 맛있는 사탕을 의미하며, 1,000,000은 가장 맛없는 사탕을 의미한다. 수정이는 동생이 말을 잘 들은 정도에 따라서, 사탕상자 안에 있는 사탕들 중 몇 번째로 맛있는 사탕을 꺼내주곤 한다. 예를 들어 말을 매우 잘 들었을 때에는 사탕상자에서 가장 맛있는 사탕을 꺼내주고, 말을 조금 잘 들었을 때에는 사탕상자에서 여섯 번째로 맛있는 사탕을 꺼내주는 식이다.

수정이가 보관하고 있는 사탕은 매우 많기 때문에 매번 사탕상자를 뒤져서 꺼낼 사탕을 골라내는 일은 매우 어렵다. 수정이를 도와주는 프로그램을 작성하시오.

▸ 입력

첫째 줄에 수정이가 사탕상자에 손을 댄 횟수 n(1 ≤ n ≤ 100,000)이 주어진다. 다음 n개의 줄에는 두 정수 A, B, 혹은 세 정수 A, B, C가 주어진다. A가 1인 경우는 사탕상자에서 사탕을 꺼내는 경우이다. 이때에는 한 정수만 주어지며, B는 꺼낼 사탕의 순위를 의미한다. 이 경우 사탕상자에서 한 개의 사탕이 꺼내지게 된다. 또, A가 2인 경우는 사탕을 넣는 경우이다. 이때에는 두 정수가 주어지는데, B는 넣을 사탕의 맛을 나타내는 정수이고 C는 그러한 사탕의 개수이다. C가 양수일 경우에는 사탕을 넣는 경우이고, 음수일 경우에는 빼는 경우이다. 맨 처음에는 빈 사탕상자에서 시작한다고 가정하며, 사탕의 총 개수는 2,000,000,000을 넘지 않는다. 또한 없는 사탕을 꺼내는 경우와 같은 잘못된 입력은 주어지지 않는다.

▸ 출력

A가 1인 모든 입력에 대해서, 꺼낼 사탕의 맛의 번호를 출력한다.

 

문제 풀이  

세그먼트의 인덱스를 사탕의 맛이라하고, 해당 인덱스에 들어있는 값을 갯수를 가진 합 세그먼트 트리로 풀이를 하면 된다. 맛이 1인 사탕 2개, 2인 사탕 1개, 3인 사탕 1개가 있다고 하면 세그먼트 트리를 다음과 같이 생성된다.

 

합 세그먼트 트리

 

B번맛 사탕 출력하기

그러면 이제 2번 쿼리만 해결해주면 된다. 현재 들어있는 사탕 중 B순위에 해당하는 것을 꺼내주는 로직으로 구간합을 이용해 이분탐색을 해주면된다. B를 target으로 한다면 다음과 같이 동작한다.

 

target = 3

• target <= tree[2] 이므로, 왼쪽 트리로 이동한다. idx = 2
• target > tree[4] 이므로, 오른쪽 트리로 이동한다. idx = 5
• 5번 인덱스에 속하는 2번맛 사탕을 출력해준다.

target = 3

 

target = 4

• target > tree[2] 이므로, 오른쪽 트리로 이동한다. idx = 3
• 3번 인덱스에 속하는 3번맛 사탕을 출력해준다.

target = 4

 

 

이를 코드로 구현하면 다음과 같다. 마지막에 캔디를 찾으면 수동으로 s번맛 사탕을 -1 감소시켜줘야 한다.

static int query(int s, int e, int idx, int target) {
	if(s == e) {
		update(1, SIZE, 1, s, -1);
		return s;
	}
	
	int mid = (s+e)/2;
	if(target <= tree[idx*2]) {
		return query(s, mid, idx*2, target);
	}else {
		return query(mid+1, e, idx*2+1, target-tree[idx*2]);
	}
}

 

풀이 코드 

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

public class Main {

	static final int SIZE = 1_000_001;
	static int[] tree;
	public static void main(String[] args) throws IOException {
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		int n = Integer.parseInt(br.readLine());
		tree = new int[SIZE*4];
		StringTokenizer st;
		StringBuilder sb = new StringBuilder();
		for(int i=0; i<n; i++) {
			st = new StringTokenizer(br.readLine());
			int op = Integer.parseInt(st.nextToken());
			int a = Integer.parseInt(st.nextToken());
			if(op == 1) {
				int candy = query(1, SIZE, 1, a);
				sb.append(candy+"\n");
			}else {
				int b = Integer.parseInt(st.nextToken());	
				update(1, SIZE, 1, a, b);
			}
		}
		System.out.println(sb.toString());
	}
	static int query(int s, int e, int idx, int target) {
		if(s == e) {
			update(1, SIZE, 1, s, -1);
			return s;
		}
		
		int mid = (s+e)/2;
		if(target <= tree[idx*2]) {
			return query(s, mid, idx*2, target);
		}else {
			return query(mid+1, e, idx*2+1, target-tree[idx*2]);
		}
	}
	
	static void update(int s, int e, int idx, int target, int dif) {
		if(target < s || e < target) return;
		
		tree[idx] += dif;
		if(s == e) return;
		
		int mid = (s+e)/2;
		update(s, mid, idx*2, target, dif);
		update(mid+1, e, idx*2+1, target, dif);
	}
}