160. Segment Tree - kth

세그먼트 트리에서 $k$번째 원소를 빠르게 찾는 방법에 대해 알아본다.

---

다음과 같은 쿼리가 $q$개 주어지는 문제를 생각해 보자.

• $1\ k$: 집합에 원소 $k$를 추가한다. $k$는 $1$ 이상 $n$ 이하의 정수 중 하나이다.
• $2\ k$: 집합에서 $k$번째로 작은 원소를 출력한다.

이 문제를 해결하는 가장 간단한 방법은 $1$번 쿼리에서 추가되는 원소를 정렬된 집합에 끼워넣는 방법과 $2$번 쿼리마다 집합을 정렬하는 방법이지만 두 방법 모두 쿼리 하나를 처리하는 시간이 $\Omega(n)$이 돼서 전체 시간복잡도는 $\Omega(nq)$로 표기할 수 있다. 물론 $n$과 $q$가 크면 속도가 너무 느려진다.

---

세그먼트 트리는 이 문제를 $\Theta(q \lg n)$에 해결하는 방법을 제공한다. 처음에 모든 원소가 $0$으로 초기화된 합 세그먼트 트리가 있다고 가정하고, $1$번 쿼리가 들어오면 $k$번째 원소의 값을 $1$ 증가시킨다. $2$번 쿼리가 들어오면 세그먼트 트리를 다음과 같이 탐색해서 답을 구할 수 있다.

• 루트에서 탐색을 시작한다. 남은 원소의 수 $x$는 $k$로 초기화한다. 만약 처음부터 집합의 크기가 $k$보다 작으면 $k$번째로 작은 원소는 존재하지 않으므로 탐색을 진행할 필요가 없다.
• 현재 노드의 왼쪽 자식 노드의 값 $y$가 $x$ 이상인 경우 왼쪽 자식 노드로 이동한다. 그렇지 않은 경우 $x$를 $x - y$로 바꾸고 오른쪽 자식 노드로 이동한다. 이 과정을 리프 노드에 도달할 때까지 반복한다.
• 리프 노드에 도달한 경우 탐색을 종료한다. 해당 리프 노드의 순서가 집합에서 $k$번째로 작은 원소와 같다.

예시를 통해 자세히 이해해 보자. 집합 $S$에서 $3$번째로 작은 원소를 찾는 것이다.

$$S = \{1, 2, 4, 6, 7\}$$

먼저 세그먼트 트리를 구성한다. 집합 $S$를 나타내는 세그먼트 트리는 다음과 같다.

여기서 $3$번째로 작은 원소를 찾아야 한다. 앞에서 설명한 방식대로 루트 노드부터 탐색을 진행한다. 처음에 $x = 3$이다.

현재 노드의 왼쪽 자식 노드의 값은 $3$이다. $3$은 $x$ 이상이므로 왼쪽 자식 노드로 이동한다. $x$의 값은 변하지 않는다.

현재 노드의 왼쪽 자식 노드의 값은 $2$이다. $2$는 $x$ 미만이므로 $x$에서 $2$를 빼고 오른쪽 자식 노드로 이동한다. $x$의 값은 $1$로 바뀐다.

현재 노드의 왼쪽 자식 노드의 값은 $0$이다. $0$은 $x$ 미만이므로 $x$에서 $0$을 빼고 오른쪽 자식 노드로 이동한다.

리프 노드에 도달했으므로 탐색을 종료한다. 또한 $4$번째 리프 노드에 도달했으므로 집합에서 $3$번째로 작은 원소가 $4$라는 것을 알 수 있다.

 

이 방법을 흔히 세그먼트 트리에서의 이분 탐색(Binary Search on Segment Tree)이라고 부른다. (다만 이름이 길어서 그대로 부르는 경우는 많지 않고 보통 줄여서 부르는 것 같다.) 세그먼트 트리를 응용할 때 꽤 자주 쓰이는 기법이므로 익숙해지는 것이 좋다. 구현은 다음과 같으며 기본적인 전처리는 기본 세그먼트 트리를 설명한 글에 나온 대로 하면 된다.

int getkth(int k, int d)
{
    if(d >= p)return d - p;
    if(a[d << 1] >= k)return getkth(k, d << 1);
    else return getkth(k - a[d << 1], d << 1 | 1);
}

호출할 때는 $d = 1$로 두면 된다. $\text{getkth}$ 함수는 비재귀로 구현할 수도 있는데, 그러면 다음과 같은 코드가 된다.

int getkth(int k)
{
    for(int d = 1; d < p;)
    {
        if(a[d << 1] >= k)d = d << 1;
        else
        {
            k -= a[d << 1];
            d = d << 1 | 1;
        }
    }
    return d - p;
}

 

[연습문제]

 

BOJ 2243. 사탕상자 (Platinum V)

세그먼트 트리에서의 이분 탐색을 사용해서 해결할 수 있는 전형적인 문제이다. 같은 원소가 여러 개 존재할 수 있고 집합에서 원소를 빼는 쿼리도 있지만 이것들은 기본적인 쿼리에서 조금만 변형하면 쉽게 처리할 수 있다.

 

BOJ 1572. 중앙값 (Platinum V)

세그먼트 트리에서의 이분 탐색을 사용해서 각각의 연속 구간의 중앙값을 빠르게 구한 뒤 다 더하면 된다.

 

BOJ 12899. 데이터 구조 (Platinum IV)

$k$번째 원소를 찾는 쿼리와 $k$번째 원소를 삭제하는 쿼리가 묶여 있는데 풀이는 위의 문제들과 크게 다르지 않다.

 

BOJ 1168. 요세푸스 문제 2 (Platinum III)

세그먼트 트리를 이용해서 각각의 위치를 기준으로 다음 $k$번째 원소를 빠르게 찾는다.