C++로 이진 트리의 대각선 합 계산하기

이진 트리에서 대각선 합(Diagonal Sum)을 구하려면 기울기가 -1인 직선들을 기준으로 노드들을 살펴봐야 합니다. 즉, 각 기준선 사이에 위치한 모든 노드의 데이터 값을 더한 것이 곧 해당 대각선의 합이 됩니다.

트리 노드 구조체 정의

먼저 노드의 데이터와 왼쪽·오른쪽 자식 노드를 담고 있는 트리 노드를 표현할 구조체를 정의합니다. 가장 처음 생성되는 노드는 루트(root) 노드가 되고, 그 이후에 생성되는 노드들은 자식(child) 노드가 됩니다.

struct Node {
    int data;
    struct Node *leftChild, *rightChild;
};

노드 생성 함수 만들기

다음으로 createNode(int data) 함수를 작성합니다. 이 함수는 int 값을 인자로 받아 새로운 노드를 생성한 뒤, 해당 값을 노드의 data 멤버에 할당하고 생성된 노드의 포인터를 반환합니다.

Node * createNode(int data){
    Node * node = new Node;
    node->data = data;
    return node;
}

대각선 합을 재귀적으로 계산하기

diagonal_sum(Node *root, int depth, map<int, int> &diagonalSum) 함수는 루트 노드와 현재 깊이(depth), 그리고 참조로 전달된 diagonalSum 맵을 인자로 받습니다. 루트가 NULL이 아니라면 현재 노드의 데이터를 diagonalSum 맵의 현재 깊이 인덱스에 더하여 요소들의 합을 누적합니다. 이후 중위 순회(inorder traversal) 방식으로 트리를 재귀적으로 탐색하며, 왼쪽 자식으로 이동할 때마다 깊이를 1씩 증가시킵니다.

void diagonal_sum(Node *root, int depth, map<int, int> &diagonalSum){
    if(root){
        diagonalSum[depth]+=root->data;
        diagonal_sum(root->leftChild, depth+1, diagonalSum);
        diagonal_sum(root->rightChild, depth, diagonalSum);
    }
}

main 함수에서 트리 구성하기

main 함수 안에서는 createNode(data) 메서드를 사용해 트리를 만들고, 합계를 저장할 맵 sumMap을 선언합니다. 루트 노드, 초기 깊이 값 1, 그리고 sumMap을 diagonal_sum 함수에 전달하면 sumMap에 키-값 쌍이 채워집니다. 이후 sumMap을 순회하기 위한 반복자(iterator) it를 생성합니다.

int main(){
    Node *root = createNode(1);
    root->rightChild = createNode(3);
    root->rightChild->leftChild = createNode(4);
    root->rightChild->leftChild->leftChild = createNode(12);
    root->rightChild->leftChild->rightChild = createNode(7);
    root->leftChild = createNode(2);
    root->leftChild->leftChild = createNode(9);
    root->leftChild->rightChild = createNode(6);
    root->leftChild->leftChild->rightChild = createNode(10);
    root->leftChild->rightChild->leftChild = createNode(11);
    root->rightChild->rightChild = createNode(5);
    map<int,int> sumMap;
    diagonal_sum(root, 1, sumMap);
    map<int,int>::iterator it;

결과 출력하기

마지막으로 for 루프 안에서 반복자 it를 사용해 sumMap을 순회하면서 각 대각선의 합을 출력합니다.

for(it=sumMap.begin(); it!=sumMap.end();++it){
    int value = it->second;
    cout<<value<<"\t";
}

전체 예제 코드

아래는 이진 트리의 대각선 합을 구하는 전체 구현 예제입니다.

#include<iostream>
#include<map>
using namespace std;
struct Node{
    int data;
    struct Node* leftChild, *rightChild;
};
Node * createNode(int data){
    Node * node = new Node;
    node->data = data;
    return node;
}
void diagonal_sum(Node *root, int depth, map<int, int> &diagonalSum){
    if(root){
        diagonalSum[depth]+=root->data;
        diagonal_sum(root->leftChild, depth+1, diagonalSum);
        diagonal_sum(root->rightChild, depth, diagonalSum);
    }
}
int main(){
    Node *root = createNode(1);
    root->rightChild = createNode(3);
    root->rightChild->leftChild = createNode(4);
    root->rightChild->leftChild->leftChild = createNode(12);
    root->rightChild->leftChild->rightChild = createNode(7);
    root->leftChild = createNode(2);
    root->leftChild->leftChild = createNode(9);
    root->leftChild->rightChild = createNode(6);
    root->leftChild->leftChild->rightChild = createNode(10);
    root->leftChild->rightChild->leftChild = createNode(11);
    root->rightChild->rightChild = createNode(5);
    map<int,int> sumMap;
    diagonal_sum(root, 1, sumMap);
    map<int,int>::iterator it;
    for(it=sumMap.begin(); it!=sumMap.end();++it){
        int value = it->second;
        cout<<value<<"\t";
    }
    return 0;
}

실행 결과

위 코드를 실행하면 다음과 같은 출력 결과를 얻을 수 있습니다.

9	19	4	2