[BOJ 12899] - 데이터 구조 (Platinum 4)

문제

자연수를 저장하는 데이터베이스 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의 쿼리 개수만큼의 줄에 각 쿼리에 대한 답을 출력합니다.

접근 방법

$Query$ $1$은 기본적으로 $O(1)$의 시간복잡도를 가진다.
하지만 $Query$ $2$를 S의 오름차순을 따라 탐색을 하는 매우 단순한 방법으로 처리를 하면 최악의 경우 $O(N)$의 시간복잡도를 가지게 된다.
즉 문제 전체의 시간복잡도는 $O(N^2)$이 되어 주어진 시간을 초과하게 된다.

처음 아이디어를 얻게된 곳은 추가되는 수인 X의 범위가 $(1 ≤ X ≤ 2,000,000)$ 로 꽤 좁다는 것이였다.
사실 잘 생각을 해보면 S에 들어있는 수 중 X번째로 작은 수를 찾는 것은 S에 들어있는 각 수의 개수를 그 수의 인덱스에 저장해 둔 배열의 1 ~ ?까지의 구간합을 이용해서 찾을 수 있다.
즉 배열의 값이 변화할 때 구간합을 효율적으로 구할 수 있게 해주는 자료구조인 "세그먼트 트리"를 이용하면
$Query$ $1, 2$ 모두 $O(log N)$의 시간 복잡도로 처리해줄 수 있게 된다.

구현

a = 2000000 # S에 추가될 수 있는 가장 큰 수
seg_tree = [0] * (a * 4) # 세그먼트 트리 배열

우선 배열의 정의를 $arr[i] =$ $S$에 들어있는 $i$의 개수로 정의했다.
처음에는 어떠한 수도 $S$에 들어있지 않으므로 세그먼트 트리 배열의 모든 값들을 0으로 초기화 해주었다.

def update(start, end, node, target, diff):
   if target < start or target > end:
       return

   seg_tree[node] += diff

   if start == end:
       return
   mid = (start + end) // 2
   update(start, mid, node * 2, target, diff)
   update(mid + 1, end, node * 2 + 1, target, diff)

$Query$ $1$을 처리해줄 update 함수이다.

def find(start, end, node, target): 
    seg_tree[node] -= 1

    if start == end:
        return start

    mid = (start + end) // 2
    if seg_tree[node * 2] >= target:
        return find(start, mid, node * 2, target)
    else:
        return find(mid + 1, end, node * 2 + 1, target - seg_tree[node * 2])

이 문제에서 가장 핵심이라고 할 수 있는 $Query$ $2$를 처리해주는 find 함수이다.
평범한 방법으로는 $O(N)$이 걸리는 쿼리를 완전 이진 트리의 특성을 이분탐색처럼 이용하여 $O(logN)$에 처리할 수 있게 해준다.

우선 $Query$에서 주어진 $X$를 target에 넣어준다. 만약 현재 노드의 왼쪽 자식 노드가 가지는 구간 합이 target 보다 같거나 크다면 $X$번째로 작은 수의 인덱스가 왼쪽 자식 노드에 있는 것임으로 왼쪽 노드를 탐색해준다.

그게아니라 만약, 왼쪽 자식 노드가 가지는 구간합이 target보다 작다면 $X$번째로 작은 수의 인덱스가 오른쪽 자식 노드에 있는 것임으로 오른쪽 노드를 탐색해준다.
이때 우리가 찾아야하는건 오른쪽 노드에 target에 맞는 노드를 찾는 것임으로
target - seg_tree[node * 2]으로 찾아야 하는 target 값을 갱신해준다.

계속 탐색을 진행해다 $start == end$ 인 리프노드에 도달하면 배열의 정의에 따라 리프노드가 가지는 범위를 반환해준다.

근데 그럼 seg_tree[node] -= 1 이건 뭘까?

처음에는 $X$번째로 작은 수를 지우는 작업을 update 함수로 따로 처리해줄려고 했다.
하지만 시간초과가 나서 좀 더 효율적으로 처리하는 방법을 고민해 보았는데,
잘생각해보면, find 함수가 방문하는 트리의 노드의 구간은 $X$가 포함되어 있다.
또한 방문했던 노드를 한 쿼리에서 다시 방문할 필요도 없기 때문에 방문하자 마자 각 노드의 값을 -1하면 되겠다고 생각했다.

후기

배열의 정의를 $arr[i] =$ $S$에 들어있는 $i$의 개수로 두는 것은 매우 빠르게 생각을 해냈는데
find 함수를 응용하는 것이 생각보다 어려웠다.
트리의 구조와 원리를 잘 이용한 좋은 문제라고 생각된다.