[자료구조/알고리즘] 세그먼트 트리 응용: K번째 수 찾기 (Segment Tree Application: Finding the K-th Element)
1. K번째 수 찾기
세그먼트 트리를 활용해 해결할 수 있는 흥미로운 문제 중 하나는 K번째 수를 찾는 문제다. 구간 합을 구하는 세그먼트 트리 와 카운팅 배열, 이분 탐색의 아이디어를 종합하면 세그먼트 트리에서 K번째로 작은 수를 $O(\log{N})$ 의 시간복잡도로 찾을 수 있다.
2. K번째 수 찾기 진행 과정
세그먼트 트리에서 K번째 수를 찾는 쿼리를 해결하기 위해서는 주어진 수의 범위를 다룰 수 있는 카운팅 배열을 세그먼트 트리의 원본 배열로 갖고 세그먼트 트리는 구간 합을 저장하면 된다. 카운팅 배열의 인덱스가 해당 수, 값이 등장 횟수를 의미하므로 세그먼트 트리의 각 노드의 값은 노드가 처리하는 구간에 포함된 수의 개수를 의미하게 된다.
주어진 수의 범위가 $0$ ~ $9$ 일 경우 초기 세그먼트 트리는 아래와 같이 된다.
$2$ 가 추가될 경우 카운팅 배열에서는 인덱스 $2$ 의 값이 $+1$ 이 되는 것이며, 세그먼트 트리에서는 아래와 같은 갱신 과정이 발생한다.
$5$ 가 추가될 경우 카운팅 배열에서는 인덱스 $5$ 의 값이 $+1$ 이 되는 것이며, 세그먼트 트리에서는 아래와 같은 갱신 과정이 발생한다.
$7$ 이 추가될 경우 카운팅 배열에서는 인덱스 $7$ 의 값이 $+1$ 이 되는 것이며, 세그먼트 트리에서는 아래와 같은 갱신 과정이 발생한다.
$0$, $4$, $4$ 가 추가될 경우 세그먼트 트리는 아래와 같이 된다.
이와 같은 상태의 세그먼트 트리에서 $3$ 번째로 작은 값을 찾아보자. K번째 수를 찾는 핵심은 왼쪽 자식 노드에 해당 수가 포함됐을지, 오른쪽 자식 노드에 해당 수가 포함됐을지 찾는 것이다. 루트 노트를 기준으로 왼쪽 자식 노드의 값이 $4$ 이므로 $0$ ~ $4$ 사이에 $4$ 개의 수가 존재한다는 것이다. 오른쪽 자식 노드에는 $3$ 번째 수가 존재할 수 없으므로 왼쪽 자식 노드만 탐색한다.
왼쪽 자식 노드의 값과 오른쪽 자식 노드의 값이 모두 $2$ 다. $0$ ~ $2$ 사이에 $2$ 개의 수가 있고, $3$ ~ $4$ 사이에 $2$ 개의 수가 있다는 의미이므로 $3$ 번째 수는 오른쪽 구간에 존재한다. 이러면 오른쪽 자식 노드를 탐색하면 되는데, 이때 $3$ 번째는 맨 앞부터 셌을 때의 순서이므로 오른쪽 자식 노드에서는 $1$ 번째 수를 찾아야 한다.
왼쪽 자식 노드의 값은 $0$, 오른쪽 자식 노드의 값은 $2$ 로 $3$ 인 수는 $0$ 개, $4$ 인 수는 $2$ 개다. 오른쪽 자식 노드를 탐색하며 오른쪽 자식 노드에서 $1$ 번째 수를 찾으면 된다.
리프 노드에 도달하면 해당 수가 K번째 수가 된다. 현재 \(\{0, 2, 4, 4, 5, 7\}\) 에서 $3$ 번째 수는 $4$ 인데 잘 찾은 것을 볼 수 있다.
3. K번째 수 찾기 코드
세그먼트 트리에서 이분 탐색의 아이디어를 활용해서 K번째 수를 구할 수 있다. 왼쪽 자식 노드를 탐색할 때는 K번째 순서를 그대로 가져가고, 오른쪽 자식 노드를 탐색할 때는 왼쪽 자식 노드의 포함된 갯수만큼 빼준 순서를 찾아야 하는 점만 주의해서 구현하면 된다.
1. K번째 수 찾기 [Java]
1
2
3
4
5
6
7
8
9
10
11
int kth(int node, int start, int end, int k) {
if (start == end) return start;
int mid = (start + end) / 2;
if (k <= tree[node * 2]) {
return kth(node * 2, start, mid, k);
} else {
return kth(node * 2 + 1, mid + 1, end, k - tree[node * 2]);
}
}
2. K번째 수 찾기 [C++]
1
2
3
4
5
6
7
8
9
10
11
int kth(int node, int start, int end, int k) {
if (start == end) return start;
int mid = (start + end) / 2;
if (k <= tree[node * 2]) {
return kth(node * 2, start, mid, k);
} else {
return kth(node * 2 + 1, mid + 1, end, k - tree[node * 2]);
}
}