세그먼트 트리 문제
해당 문제의 응용문제입니다. 위 문제도 같이 푸는 것을 권장합니다.
[문제 풀이]
2번 쿼리마다 k번째 수를 반환하는 문제입니다.
1번 쿼리는 자연수를 추가하는 역할입니다.
쿼리의 개수가 심상치 않습니다.
쿼리의 개수를 보니 매 쿼리를 logn정도로 처리해야 시간 초과가 걸리지 않을 것 같습니다.
세그먼트 트리로 접근해 봅시다.
이 문제를 처음 세그먼트 트리로 생각해보면 상당히 막막합니다.
약간의 팁을 드리면 대부분의 세그먼트 트리 문제는 구간 합, 최소, 최대를 생각하고 접근하는 것이 좋습니다.
세그먼트 트리라고 생각은 들지만 평소에 풀던 세그먼트 트리와는 조금 다른 형태입니다.
주어진 수열도 없고, 그에 따라 세그먼트 트리의 크기도 잡기가 상당히 막막합니다.
하지만 입력 조건을 잘 보면 X의 범위가 2,000,000 이하의 자연수입니다.
인덱스로 쓰기에 너무 좋은 범위이고 [1, 2,000,000] 범위의 구간 합 세그먼트 트리를 만들어서 이용한다면 그럴듯하게 구할 수 있을 것 같습니다.
더구나 1번 쿼리는 세그먼트 트리의 update과정과 매우 흡사합니다.
문제는 어떻게 k번째 수를 구하는가에 대한 아이디어와 어떤 식으로 구간 합 세그먼트 트리를 이용하는지에 대한 것입니다.
메인 아이디어는 아래와 같습니다.
- 1번 쿼리는 X에 해당하는 모든 노드에 +1을 하는 것
- 2번 쿼리는 k번째 수 X를 찾아서 출력하고 X에 해당하는 모든 노드에 -1을 하는 것
k번째 수를 찾는 방법은 이분 탐색과 상당히 흡사한데, 현재 노드에서 범위를 반으로 나누는 자식 노드가 두 개가 있으므로 구간을 반씩 계속 나눠가면서 자식 노드의 구간 합 (원소의 개수)이 k를 넘지 않도록 하면 구할 수 있습니다.
고민은 많이 했으나 쉽게 설명하지 못한 것 같아 문제를 푸는 로직을 간단히 설명하자면
- 1번 쿼리에서는 구간에 대해 +1 합니다.
- 2번 쿼리에서는 k보다 왼쪽 자식 노드의 구간 합이 작다면 k - (왼쪽 자식 노드의 구간 합)을 하고 오른쪽 자식으로 이동하고, k보다 왼쪽 자식 노드의 구간 합이 크거나 같다면 왼쪽 자식으로 이동합니다. 그리고 얻은 k번째 수의 구간에 대해 -1
그림으로 설명하는 것이 이해하기 훨씬 편하므로 예시를 들어서 그림으로 표현하겠습니다.
해당 예시에서는 세그먼트 트리의 범위가 [1, 5]이고
5
1 1
1 2
1 5
2 2
2 2
와 같은 입력이 들어왔다고 가정하겠습니다.
5
1 1
1 2
1 5
까지 입력이 들어왔다면 세그먼트 트리는 아래와 같이 구성이 되어 있을 겁니다.
그리고 2 2라는 입력이 추가로 들어왔다고 가정하면
들어온 k값 2를 현재 노드인 [1, 5]의 왼쪽 노드인 [1, 3]의 구간 합과 비교하게 됩니다.
k와 [1, 3]의 구간 합 2는 같으므로 right를 줄여줍시다.
right를 줄였으므로 현재 노드는 [1, 3]을 가리키게 되고 k를 왼쪽 자식 노드인 [1, 2]의 구간 합과 비교하게 됩니다.
[1, 2]도 2이고 k와 같으므로 또 right를 줄여줍시다.
right를 줄였으므로 현재 노드는 [1, 2]를 가리키게 되고 k를 왼쪽 자식 노드인 1의 구간 합과 비교하게 됩니다.
1의 구간 합은 1이므로 첫번째 수가 1 임을 알 수 있습니다.
k가 2이므로 우리는 두번째 수를 구하는 것이고 이번에도 right를 줄이게 된다면 첫 번째 수인 1만 남게 되므로 이번에는 left를 올려줍시다.
left == right == 2가 되었고 2가 두번째 수임을 알 수 있게 되었습니다.
두 번째 수인 2를 알게 되었으므로 출력하고 세그먼트 트리에서 2가 포함되는 모든 부분에 대해 -1을 더해줍시다.
그렇다면 위의 트리처럼 변합니다.
다시 2 2의 입력이 추가로 들어왔다고 생각합시다.
들어온 k값 2를 현재 노드의 왼쪽 자식 노드의 구간 합 1과 비교합니다.
[1, 3]의 구간 합이 1이므로 [1, 3]의 범위에는 수가 한 개 뿐입니다.
[1, 3]에 한 개의 수가 있음을 확인했으므로 두 번째 수를 찾기 위해서는 left를 올려주고 남은 범위에서의 첫 번째 수를 찾으면 됩니다.
더 간단히 얘기하자면 k > (왼쪽 노드의 구간 합) 이면 오른쪽 노드의 구간에서 k - (왼쪽 노드의 구간 합) 만큼을 다시 탐색하면 됩니다.
현재 k가 1이고 (2 - 1) 왼쪽 노드의 구간 합은 0이므로 이번에도 left를 늘려줍시다.
달라진 점은 없지만 일반성을 유지하기 위해 k - 0도 해줍니다.
left == right == 5가 되었고 5가 두 번째 수임을 알 수 있게 되었습니다.
위와 같은 방식으로 $O(Qlogn)$의 시간 복잡도로 모든 쿼리를 처리할 수 있습니다.
[소스 코드] / Java
import java.io.*;
import java.util.*;
public class Main {
private static final int SIZE = 2000000;
private static int[] sumtree = new int[SIZE * 4];
public static void main(String[] args) throws IOException {
BufferedReader br = new BufferedReader(new InputStreamReader(System.in));
int N = Integer.parseInt(br.readLine());
StringBuilder sb = new StringBuilder();
StringTokenizer st;
for (int Q = 0; Q < N; Q++) {
st = new StringTokenizer(br.readLine(), " ");
char q = st.nextToken().charAt(0);
int i = Integer.parseInt(st.nextToken());
if (q == '1') {
update(1, 6, 1, i, 1);
} else {
int n = search(1, 6, 1, i);
update(1, 6, 1, n, -1);
sb.append(n).append('\n');
}
}
System.out.println(sb.toString());
br.close();
}
private static int search(int left, int right, int node, int i) {
while (left < right) {
node *= 2;
int l = sumtree[node], mid = (left + right) / 2;
if (l < i) {
left = mid + 1;
node++;
i -= l;
} else {
right = mid;
}
}
return left;
}
private static int update(int left, int right, int node, int i, int k) {
if (left <= i && i <= right) {
if (left == right) return sumtree[node] += k;
int mid = (left + right) / 2;
sumtree[node] = update(left, mid, node * 2, i, k) + update(mid + 1, right, node * 2 + 1, i, k);
}
return sumtree[node];
}
}
'Algorithm > BOJ' 카테고리의 다른 글
| [백준 20152] Game Addiction (0) | 2021.03.24 |
|---|---|
| [백준 7662] 이중 우선순위 큐 (0) | 2021.03.10 |
| [백준 11501] 주식 (0) | 2021.01.25 |
| [백준 20500] Ezreal 여눈부터 가네 ㅈㅈ (0) | 2021.01.24 |
| [백준 20149] 선분 교차 3 (0) | 2021.01.21 |
