[BOJ] 백준 12899번 데이터 구조 (Java)

#12899 데이터 구조

난이도 : 플레 4

유형 : 세그먼트 트리 / 인덱스 트리

12899번: 데이터 구조

▸ 문제

자연수를 저장하는 데이터베이스 S에 대해 다음의 쿼리를 처리합시다.

유형 1 : S에 자연수 X를 추가한다.

유형 2 : S에 포함된 숫자 중 X번째로 작은 수를 응답하고 그 수를 삭제한다.

▸ 입력

첫째 줄에 사전에 있는 쿼리의 수 N 이 주어집니다. (1 ≤ N ≤ 2,000,000)

둘째 줄부터 N개의 줄에 걸쳐 각 쿼리를 나타내는 2개의 정수 T X가 주어집니다.

T가 1이라면 S에 추가할 X가 주어지는 것입니다. (1 ≤ X ≤ 2,000,000)

T가 2라면 X는 S에서 삭제해야 할 몇 번째로 작은 수인지를 나타냅니다. S에 최소 X개의 원소가 있음이 보장됩니다.

▸ 출력

유형 2의 쿼리 개수만큼의 줄에 각 쿼리에 대한 답을 출력합니다.

 

문제 풀이  

최대 쿼리의 수는 200만이고 원소의 개수도 최대 200만이다. x번째 수를 찾기 위해서 선형탐색으로 이뤄지면 시간초과로 찾을 수가 없다. 이럴 때는 트리의 구조를 사용해서 탐색의 속도를 로그함수로 줄여줘야 한다. k번째의 수는 인덱스 트리를 사용하면 간단하게 구할 수 있다.

 

인덱스 트리에 데이터 저장하기

인덱스 트리는 리프 노드에 데이터를 저장하고 합 트리로 트리 구조를 형성하는 트리이다. 예를들어 크기가 4인 배열이 있다고 하자. 그러면 인덱스 트리는 다음과 같이 생성할 수 있다.

인덱스 트리 생성하기

 

여기서 2를 추가하면 다음과 같이 카운트를 해줄 수 있다.

 

값 추가하기

 

K번째 수 찾기

배열의 데이터를 인덱스 트리로 표현한 다음 K번째 수는 다음과 같이 찾아주면 된다. 만약 [1, 3, 4]라는 값이 들어있다고 하자. 노드의 값은 데이터 개수를 나타내므로 K번째에 해당하는 구간을 찾아주면 된다.

• tree[2*node] < k :  왼쪽 자식 노드에 포함된 데이터 개수가 k개보다 작을 경우 오른쪽 자식 노드 탐색 
  • 왼쪽 노드에 들어있는 개수는 카운트한 것이므로 k = k-tree[2*node]로 갱신해준다.
• tree[2*node] >= k : k보다 많거나 같다면 해당 서브트리에 k번째 수가 들어있으므로 왼쪽 자식 노드를 탐색한다.

 

K번째 수 찾기

 

 

풀이 코드 

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

public class Main {
	static int[] arr, tree;
	static final int MAX = 2_000_001;
	public static void main(String[] args) throws IOException{
		BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
		int q = Integer.parseInt(br.readLine());

		StringBuilder sb = new StringBuilder();
		tree = new int[MAX*4];
		StringTokenizer st;
		while(q-- > 0) {
			st = new StringTokenizer(br.readLine());
			int op = Integer.parseInt(st.nextToken());
			int x = Integer.parseInt(st.nextToken());
			if(op == 1) {
				update(1, MAX, 1, x);
			}else {
				sb.append(query(1, MAX, 1, x)+"\n");
			}
		}
		System.out.println(sb.toString());
	}
	
	static void update(int s, int e, int node, int idx) {
		if(idx < s|| e < idx) return;
		tree[node] += 1;
		if(s == e) return;
		int mid = (s+e)/2;
		update(s, mid, node*2, idx);
		update(mid+1, e, node*2+1, idx);
	}
	
	static int query(int s, int e, int node, int k) {
		tree[node]-=1;
		if(s == e) {
			return s;
		}
		
		int mid = (s+e)/2;
		if(tree[2*node] < k) { // 왼쪽 자식 노드에 포함된 데이터 개수가 k개보다 작을 경우 오른쪽 탐색    
			return query(mid+1, e, 2*node+1, k - tree[2*node]); // k-tree[2*node]번째 찾기 
		}else { // 2개와 같거나 많을 경우  
			return query(s, mid, 2*node, k);
		}
	}

}