세그먼트 트리로 k번째 작은 수 찾기

이 게시글에서는 길이 N의 수열에서 k번째로 작은 수를 세그먼트 트리로 구할 것이다.

 

질의당 $O(\log N)$의 시간복잡도가 요구된다.

 

세그먼트 트리를 구성할 때 배열의 원소값을 index로 하여 해당하는 수가 몇 개 존재하는지 저장한다.

 

그러면 update 연산은 구간합을 구하는 일반적인 세그먼트 트리의 그것과 동일한 구현을 쓸 수 있다.

 

k번째로 작은 수를 찾는 연산을 할 경우는 k를 매개변수로 전달해줘야 한다.

 

전달받은 k를 가지고 두 개의 분기로 가지치기를 할 수 있다.

 

1. 왼쪽 구간에 존재하는 원소의 개수가 k보다 크거나 같으면 왼쪽 구간으로 탐색을 한다.

2. 1에 해당하지 않으면 오른쪽 구간으로 탐색한다.

 

2번에서 오른쪽 구간으로 가지치기를 하게 되면 왼쪽 구간의 원소들을 배제하게 된다. 이때는 인자로 전달되는 k에서 왼쪽 구간에 존재하는 원소의 개수만큼 빼야 한다.

 

이렇게 start와 end가 같아질 때까지 탐색하여 얻어지는 index가 바로 k번째로 작은 원소가 되는 것이다.

이 질의는 다음과 같은 함수로 작성할 수 있게 된다.

 

int query(vector<int> &tree, int node, int start, int end, int kth)
{
    if (start == end) return end;
    int m = mid(start, end);
    int left_val = tree[node * 2];

    if (kth <= left_val) return query(tree, node * 2, start, m, kth);
    return query(tree, node * 2 + 1, m + 1, end, kth - left_val);
}

 

왼쪽 구간에 해당하는 left_val을 가져와 탐색 분기를 결정하고 start와 end가 같아질 때까지 탐색하는 것을 볼 수 있다.

 

이 함수를 이용하여 k번째 작은 수를 쉽게 찾을 수 있을 것이다. 관련 문제를 풀어보자.

 

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

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