이 문제에서는 하나의 이진 트리(binary tree)와 Q개의 쿼리가 주어집니다. 우리의 과제는 이진 트리의 두 노드 사이의 거리를 찾는 쿼리를 효율적으로 처리하는 프로그램을 작성하는 것이며, 이 글에서는 C++로 O(log n) 시간 복잡도를 달성하는 방법을 소개합니다.
문제 설명
각 쿼리마다 이진 트리의 두 노드가 주어지며, 두 노드 사이의 거리, 즉 한 노드에서 다른 노드로 이동하기 위해 거쳐야 하는 간선(edge)의 개수를 구해야 합니다.
예시를 통해 문제를 살펴보겠습니다.
입력: 이진 트리
쿼리 수 = 3
Q1 -> [2, 6]
Q2 -> [4, 1]
Q3 -> [5, 3]
출력: 3, 2, 3
풀이 접근 방식
이 문제는 최소 공통 조상(Lowest Common Ancestor, LCA)과 루트에서 각 노드까지의 거리를 활용한 거리 공식으로 해결할 수 있습니다.
Distance(n1, n2) = distance(root, n1) + distance(root, n2) - 2 * distance(root, LCA)
문제를 해결하기 위해 다음 단계를 따릅니다.
- 각 노드(N1, N2, LCA)의 레벨(깊이)을 구합니다.
- 트리의 오일러 경로(Euler Tour)를 기반으로 배열을 생성합니다.
- LCA 질의를 빠르게 처리하기 위한 세그먼트 트리(segment tree)를 구축합니다.
예제 코드
#include <bits/stdc++.h>
#define MAX 1000
using namespace std;
int eulerArray[MAX];
int eIndex = 0;
int vis[MAX];
int L[MAX];
int H[MAX];
int level[MAX];
struct Node {
int data;
struct Node* left;
struct Node* right;
};
struct Node* newNode(int data) {
struct Node* temp = new struct Node;
temp->data = data;
temp->left = temp->right = NULL;
return temp;
}
void FindNodeLevels(struct Node* root) {
if (!root)
return;
queue<pair<struct Node*, int> > q;
q.push({ root, 0 });
pair<struct Node*, int> p;
while (!q.empty()) {
p = q.front();
q.pop();
level[p.first->data] = p.second;
if (p.first->left)
q.push({ p.first->left, p.second + 1 });
if (p.first->right)
q.push({ p.first->right, p.second + 1 });
}
}
void createEulerTree(struct Node* root) {
eulerArray[++eIndex] = root->data;
if (root->left) {
createEulerTree(root->left);
eulerArray[++eIndex] = root->data;
}
if (root->right) {
createEulerTree(root->right);
eulerArray[++eIndex] = root->data;
}
}
void creareEulerArray(int size) {
for (int i = 1; i <= size; i++) {
L[i] = level[eulerArray[i]];
if (vis[eulerArray[i]] == 0) {
H[eulerArray[i]] = i;
vis[eulerArray[i]] = 1;
}
}
}
pair<int, int> seg[4 * MAX];
pair<int, int> min(pair<int, int> a, pair<int, int> b) {
if (a.first <= b.first)
return a;
else
return b;
}
pair<int, int> buildSegTree(int low, int high, int pos) {
if (low == high) {
seg[pos].first = L[low];
seg[pos].second = low;
return seg[pos];
}
int mid = low + (high - low) / 2;
buildSegTree(low, mid, 2 * pos);
buildSegTree(mid + 1, high, 2 * pos + 1);
seg[pos] = min(seg[2 * pos], seg[2 * pos + 1]);
}
pair<int, int> LCA(int qlow, int qhigh, int low, int high, int pos) {
if (qlow <= low && qhigh >= high)
return seg[pos];
if (qlow > high || qhigh < low)
return { INT_MAX, 0 };
int mid = low + (high - low) / 2;
return min(LCA(qlow, qhigh, low, mid, 2 * pos), LCA(qlow, qhigh,mid + 1, high, 2 * pos +1));
}
int CalcNodeDistance(int node1, int node2, int size) {
int prevn1 = node1, prevn2 = node2;
node1 = H[node1];
node2 = H[node2];
if (node2 < node1)
swap(node1, node2);
int lca = LCA(node1, node2, 1, size, 1).second;
lca = eulerArray[lca];
return level[prevn1] + level[prevn2] - 2 * level[lca];
}
int main() {
int N = 6;
Node* root = newNode(1);
root->left = newNode(2);
root->right = newNode(3);
root->left->left = newNode(4);
root->left->right = newNode(5);
root->right->left = newNode(6);
FindNodeLevels(root);
createEulerTree(root);
creareEulerArray(2 * N - 1);
buildSegTree(1, 2 * N - 1, 1);
int Q = 4;
int query[Q][2] = {{1, 5}, {4, 6}, {3, 4}, {2, 4} };
for(int i = 0; i < Q; i++)
cout<<"The distance between two nodes of binary tree is "<<CalcNodeDistance(query[i][0], query[i][1], 2 * N - 1)<<endl;
return 0;
}
출력 결과
The distance between two nodes of binary tree is 2 The distance between two nodes of binary tree is 4 The distance between two nodes of binary tree is 3 The distance between two nodes of binary tree is 1
복잡도 분석
전처리 단계(각 노드의 레벨 계산, 오일러 경로 배열 생성, 세그먼트 트리 구축)에는 O(n)의 시간이 소요됩니다. 이후 각 거리 쿼리는 세그먼트 트리의 범위 최솟값(range minimum query) 탐색 덕분에 O(log n) 안에 처리되므로, 전체 시간 복잡도는 전처리 O(n) + Q × O(log n)입니다. 이를 통해 노드 수가 많은 트리에서도 대량의 쿼리를 효율적으로 처리할 수 있습니다.