세그먼트 트리 - (3). K번째 원소 찾기

이번 글은 세그먼트 트리 활용중 하나인 K번째 원소 찾기 입니다.

 

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

12899번: 데이터 구조

 

문제에는 두 가지 종류의 쿼리가 존재한다.

유형 1 : 데이터베이스 S에 자연수 X를 추가한다.

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

 

쿼리는 최대 200만개 까지 가능하니 한 쿼리를 최대 O(NlogN)에 처리해야 한다.

다행히도(?) X의 범위는 1부터 200만 까지이다.

 

이 문제는 X의 범위만큼 세그먼트 트리를 가져간다.

즉, leaf 노드의 갯수가 200만이 되도록 구간합 세그먼트 트리를 생성한다.

 

유형 1이 들어왔을 때 X번째 노드에 1을 추가하고(+) 부모로 올라가 다시 구간합을 구하는 방식으로 업데이트를 진행한다.

 

문제는 유형 2의 쿼리인데,

답은 생각보다 간단하다.

 

현재 구간합 노드에서 왼쪽 자식보다 X값이 크다면, X값에 왼쪽 자식의 값을 빼고 오른쪽 자식 쪽으로 탐색을 한다.

반대로 현재 구간합 노드에서 왼쪽 자식보다 X값이 작거나 같다면 왼쪽 자식 쪽으로 탐색을 한다.

 

위 그림에서 맨 마지막의 검의 글씨가 해당 자연수의 값이고 그 위의 리프 노드가 그 자연수의 갯수를 뜻한다.

"만약 7번째로 작은수를 찾아라." 라는 쿼리를 받으면 루트노드에서 왼쪽 자식에는 6개의 수밖에 없음을 알 수 있다.

그렇다면 오른쪽 자식의 자식들 중에 있다는 얘기가된다. 대신 오른쪽 자식의 자식들 중에서 (7-6)1번째로 작은 수가 된다.

그럼 5가 루트인 서브 트리에서 다시 왼쪽 자식과 오른쪽 자식을 보면 된다.

여기서 1번째로 작은 수를 찾아야 하니, 왼쪽 자식의 3보다 작은 1번째 수이므로, 왼쪽 자식의 서브트리중에 우리가 찾는 수가 있다는 얘기가 된다. 그렇게 왼쪽 자식으로 탐색을 들어가서 결국 우리는 5라는 수를 찾게된다.

 

 

이 방법으로 K번째 수를 찾으려면 다음 조건을 만족해야한다.

주어지는 수의 범위를 알고, 그 크기만큼 세그먼트 트리를 만들어야한다.

주어지는 수가 N일 때, 세그먼트 트리의 크기를 대충 4N이라고 하면, 해당 크기만큼 세그먼트 트리를 만들었을 때 메모리 제한 안에 들어와야한다.

 

이 문제의 풀이 코드는 다음과 같다.

 

#include <bits/stdc++.h>
 
using namespace std;
 
#define swap(a,b) (a)^=(b)^=(a)^=(b)
#define endl '\n'
#define int long long
 
using pii = pair<int, int>;
 
const double EPS = 1e-9;
const int dx[] = {1, -1, 0, 0};
const int dy[] = {0, 0, 1, -1};
 
vector<int> seg_tree;
 
int find(int L, int R, int value, int nodeNum)
{
    if(L == R) return L;
    int mid = (L+R)/2;
    if(value <= seg_tree[nodeNum*2]) return find(L, mid, value, nodeNum*2);
    return find(mid+1, R, value-seg_tree[nodeNum*2], nodeNum*2+1);
}
 
void update(int nodeNum, int value)
{
    seg_tree[nodeNum] += value;
    while(nodeNum > 1)
    {
        nodeNum /= 2;
        seg_tree[nodeNum] = seg_tree[nodeNum*2] + seg_tree[nodeNum*2+1];
    }
}
 
int32_t main()
{
    cin.sync_with_stdio(NULL);
    cin.tie(NULL);
    cout.tie(NULL);
 
    int depth = (int)ceil(log2(2000001))+1;
    int tree_size = 1 << depth;
    int leaf_size = tree_size / 2;
 
    seg_tree.resize(tree_size, 0);
 
 
    int n;
    cin >> n;
    while(n--)
    {
        int t, x;
        cin >> t >> x;
        if(t == 1)
        {
            update(x+leaf_size, 1);
        }
        else
        {
            int ans = find(0, leaf_size-1, x, 1);
            cout << ans << endl;
            update(ans+leaf_size, -1);
        }
    }
 
    return 0;
}

 

이 방법으로 다음 문제도 해결할 수 있다.

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

1168번: 요세푸스 문제 2

 

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

1572번: 중앙값