C++로 이진 트리의 두 노드 사이 거리 구하는 방법

이 문제에서는 하나의 이진 트리(Binary Tree)와 두 개의 노드가 주어집니다. 우리의 목표는 이진 트리에서 두 노드 사이의 거리(distance)를 구하는 프로그램을 작성하는 것입니다.

문제 설명

여기서 말하는 '두 노드 사이의 거리'란, 한 노드에서 다른 노드로 이동할 때 거쳐야 하는 간선(edge)의 최소 개수를 의미합니다.

예시를 통해 문제를 살펴보겠습니다.

입력: 이진 트리

Node1 = 3, Node2 = 5

출력: 3

설명

노드 3에서 노드 5까지의 경로는 3 → 1 → 2 → 5 입니다. 이 경로에서 총 3개의 간선을 지나게 되므로, 두 노드 사이의 거리는 3이 됩니다.

해결 접근 방법

이 문제를 푸는 가장 간단한 방법은 두 노드의 최소 공통 조상(LCA, Lowest Common Ancestor)을 찾은 뒤, 아래 공식을 적용하는 것입니다.

distance(node1, node2) = distance(root, node1) + distance(root, node2) − 2 × distance(root, LCA)

즉, 루트에서 각 노드까지의 거리를 더한 후, 루트에서 최소 공통 조상까지의 거리를 두 번 빼주면 두 노드 사이의 거리가 계산됩니다.

C++ 구현 예제

#include <iostream>
using namespace std;
struct Node{
   struct Node *left, *right;
   int key;
};
Node* insertNode(int key){
   Node *temp = new Node;
   temp->key = key;
   temp->left = temp->right = NULL;
   return temp;
}
int calcNodeLevel(Node *root, int val, int level) {
   if (root == NULL)
      return -1;
   if (root->key == val)
      return level;
   int lvl = calcNodeLevel(root->left, val, level+1);
   return (lvl != -1)? lvl : calcNodeLevel(root->right, val, level+1);
}
Node *findDistanceRec(Node* root, int node1, int node2, int &dist1, int &dist2, int &dist, int lvl){
   if (root == NULL) return NULL;
   if (root->key == node1){
      dist1 = lvl;
      return root;
   }
   if (root->key == node2){
      dist2 = lvl;
      return root;
   }
   Node *leftLCA = findDistanceRec(root->left, node1, node2, dist1,dist2, dist, lvl+1);
   Node *rightLCA = findDistanceRec(root->right, node1, node2, dist1,dist2, dist, lvl+1);
   if (leftLCA && rightLCA){
      dist = dist1 + dist2 - 2*lvl;
      return root;
   }
   return (leftLCA != NULL)? leftLCA: rightLCA;
}
int CalcNodeDistance(Node *root, int node1, int node2) {
   int dist1 = -1, dist2 = -1, dist;
   Node *lca = findDistanceRec(root, node1, node2, dist1, dist2, dist, 1);
   if (dist1 != -1 && dist2 != -1)
      return dist;
   if (dist1 != -1){
      dist = calcNodeLevel(lca, node2, 0);
      return dist;
   }
   if (dist2 != -1){
      dist = calcNodeLevel(lca, node1, 0);
      return dist;
   }
   return -1;
}
int main(){
   Node * root = insertNode(1);
   root->left = insertNode(2);
   root->right = insertNode(3);
   root->left->left = insertNode(4);
   root->left->right = insertNode(5);
   root->right->left = insertNode(6);
   cout<<"Distance between node with value 5 and node with value 3 is"<<CalcNodeDistance(root, 3, 5);
   return 0;
}

실행 결과

Distance between node with value 5 and node with value 3 is 3

마무리

위 알고리즘은 트리를 한 번 순회하며 두 노드의 위치와 최소 공통 조상을 동시에 찾기 때문에 시간 복잡도는 O(N)입니다. 여기서 N은 트리의 전체 노드 수입니다. 재귀 호출로 인한 스택 메모리 사용 외에 추가 공간이 필요하지 않아 공간 복잡도 역시 효율적입니다.