[코테] 백준 2243번 : 사탕상자 (java)

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

알고리즘 분류 : 자료 구조, 세그먼트 트리, 이분 탐색

 

❓문제

---

🔅해석

 

세그먼트 트리에서 구간의 정보를 조회하는 query는 총 다섯가지 정보를 넘긴다.

query(int left, int right, int node, int queryLeft, int queryRight) 

현재 문제에서는 해당 순위의 정보만 필요하므로 query(int left, int right, int node, int query) 으로 넘긴다.

 

문제의 입력을 보자면, 레벨 1 사탕 2개, 레벨 3 사탕 3개가 존재하고 여기서 2번째로 맛있는 사탕을 꺼내준다고 할 때, 답은 레벨 1이다.

 

위 그림을 볼 때, 노드를 탐색하면서 해당 노드의 값이 B보다 크거나 같을 때 왼쪽 자식을 호출한다.

리프 노드에 도달했을 때, 목적이이면 해당 노드의 left 또는 right를 반환하고 현재 노드와 부모 노드 value에 -1를 해준다.

 

그 다음으로 다시 2번째 맛있는 사탕을 꺼낸다.

노드를 탐색하면서 해당 노드의 값이 B보다 작을 때, 오른쪽 자식을 호출한다.

 

다만, 기존 세그먼트 트리에서 query를 날렸을 때, query의 범위가 계속 같았지만 해당 문제에서는 순위와 갯수를 판단해야하기에 query의 범위가 바뀐다.

조건은 현재 노드의 값이 꺼낼 사탕 순위(B) 보다 작다면 해당 노드의 값을 빼주고 넘겨야한다.

---

⭕정답 코드

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.StringTokenizer;

public class Main {

	static int N;
	static int S;
	static long[] tree;

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

		N = Integer.parseInt(br.readLine());

		S = 1;
		while (S < 1000000) {
			S *= 2;
		}
		tree = new long[S * 2];

		for (int i = 0; i < N; i++) {
			StringTokenizer st = new StringTokenizer(br.readLine());
			int A = Integer.parseInt(st.nextToken());
			if (A == 1) {
				long B = Long.parseLong(st.nextToken());
				System.out.println(query(1, S, 1, B));

			} else {
				int B = Integer.parseInt(st.nextToken());
				long C = Long.parseLong(st.nextToken());
				update(1, S, 1, B, C);
			}
		}
		br.close();
	}

	private static long query(int left, int right, int node, long rank) {

		if (left == right) {
			update(1, S, 1, left, -1);
			return left;
		}

		int mid = (left + right) / 2;
		if (tree[node * 2] >= rank) {
			return query(left, mid, node * 2, rank);
		} else {
			return query(mid + 1, right, node * 2 + 1, rank - tree[node * 2]);
		}
	}

	private static void update(int left, int right, int node, int target, long num) {
		// 영향이 없는 경우
		if (target < left || right < target) {
			return;
		} else {
			// 영향이 있는 경우
			tree[node] += num;
			if (left != right) {
				int mid = (left + right) / 2;
				update(left, mid, node * 2, target, num);
				update(mid + 1, right, node * 2 + 1, target, num);
			}
		}
	}
}

 

---

❗결과