C++에서 O(n)으로 이진 트리의 지름 구하기

이진 트리의 지름이란?

이진 트리의 지름(diameter)은 각 노드를 기준으로 (왼쪽 서브트리의 높이 + 오른쪽 서브트리의 높이 + 1)로 정의됩니다. 즉, 트리 안의 임의의 두 노드를 연결하는 경로 중 가장 긴 경로에 포함된 노드의 개수를 의미합니다.

이 방법에서는 모든 노드에 대해 (왼쪽 높이 + 오른쪽 높이 + 1) 값을 계산하고, 그중 최댓값으로 결과를 갱신합니다. 단 한 번의 순회만으로 답을 구할 수 있기 때문에 시간 복잡도는 O(n)으로 유지됩니다.

노드 구조체 정의

먼저 데이터와 왼쪽·오른쪽 자식 포인터를 가지는 트리 노드를 나타내는 구조체를 정의합니다. 처음 생성되는 노드는 루트 노드가 되고, 그 이후에 만들어지는 노드는 자식 노드가 됩니다.

struct Node {
    int data;
    Node* left;
    Node* right;
};

newNode() 함수

다음으로 int 값을 인자로 받아 노드의 data 멤버에 할당하는 newNode(int data) 함수를 작성합니다. 이 함수는 생성된 Node 구조체의 포인터를 반환하며, 새로 만들어진 노드의 왼쪽과 오른쪽 자식은 NULL로 초기화됩니다.

struct Node* newNode(int data){
    struct Node* node = new Node;
    node->data = data;
    node->left = node->right = NULL;
    return node;
}

diameter() 함수

diameter(Node* root) 함수는 루트 노드를 받아 해당 노드가 NULL인지 먼저 확인합니다. 그런 다음 INT_MIN 값으로 초기화된 ans 변수를 선언하고, height(root, ans)의 반환값을 height_of_tree 변수에 저장한 뒤 최종적으로 ans를 반환합니다.

int diameter(Node* root){
    if (root == NULL)
        return 0;
    int ans = INT_MIN;
    int height_of_tree = height(root, ans);
    return ans;
}

height() 함수

height(Node* root, int& ans) 함수는 루트 노드와 참조(reference)로 전달된 ans 변수를 받습니다. 트리를 순회하면서 각 서브트리의 높이를 계산하고, 재귀 호출마다 ans의 최댓값을 두 번째 매개변수로 함께 전달합니다. 재귀 호출이 끝나면 ans는 max(ans, 1 + left_height + right_height) 값으로 갱신되며, 이 값이 곧 트리의 지름 후보가 됩니다.

전체 구현 예제

다음은 O(n) 방법으로 이진 트리의 지름을 구하는 전체 코드입니다.

#include <iostream>
using namespace std;

struct Node {
    int data;
    Node* left;
    Node* right;
};

struct Node* newNode(int data){
    struct Node* node = new Node;
    node->data = data;
    node->left = node->right = NULL;
    return node;
}

int height(Node* root, int& ans){
    if (root == NULL)
        return 0;
    int left_height = height(root->left, ans);
    int right_height = height(root->right, ans);
    ans = max(ans, 1 + left_height + right_height);
    return 1 + max(left_height, right_height);
}

int diameter(Node* root){
    if (root == NULL)
        return 0;
    int ans = INT_MIN;
    int height_of_tree = height(root, ans);
    return ans;
}

int main(){
    struct Node* root = newNode(1);
    root->left = newNode(2);
    root->right = newNode(3);
    root->left->left = newNode(4);
    root->left->right = newNode(5);
    printf("Diameter is %d\n", diameter(root));
    return 0;
}

실행 결과

위 코드를 실행하면 다음과 같은 출력이 나옵니다.

Diameter is 4

이 예제 트리에서 가장 긴 경로는 리프 노드 4에서 노드 2, 루트 1을 거쳐 노드 3에 이르는 경로로, 총 4개의 노드를 포함하므로 지름은 4가 됩니다.

이 알고리즘은 각 노드를 정확히 한 번씩만 방문하므로 시간 복잡도는 O(n)이며, 재귀 호출 스택의 깊이가 트리의 높이에 비례하므로 공간 복잡도는 O(h)(h는 트리의 높이)입니다.