Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

C++로 이진 트리 수직 순회에서 K번째 노드 찾는 방법

문제 개요

이진 트리와 값 K가 주어졌을 때, 수직 순회(vertical order traversal) 결과에서 K번째 노드를 찾아 출력하는 것이 이 글의 목표입니다. 만약 해당 위치의 노드가 존재하지 않는다면 -1을 반환해야 합니다.

예를 들어 다음과 같은 이진 트리가 있다고 가정해 보겠습니다.

C++로 이진 트리 수직 순회에서 K번째 노드 찾는 방법

이 트리를 수직 순회하면 노드는 아래와 같은 순서로 방문됩니다.

4
2
1 5 6
3 8
7
9

위 순서에서 K = 3이라면, 세 번째 위치에 있는 노드는 1이므로 결과값은 1이 됩니다.

접근 방법

풀이 방식은 비교적 간단합니다. 먼저 너비 우선 탐색(BFS)을 기반으로 수직 순회를 수행한 뒤, 노드를 하나씩 세어가면서 현재 노드가 K번째 노드인지 확인하고, 해당 노드를 찾으면 즉시 그 값을 반환합니다.

핵심 아이디어를 단계별로 정리하면 다음과 같습니다.

  • 각 노드에 수평 거리(horizontal distance)를 부여합니다. 루트는 0, 왼쪽 자식은 부모보다 1 작게, 오른쪽 자식은 부모보다 1 크게 설정합니다.
  • 큐(queue)를 사용해 트리를 레벨 순서대로 순회하면서, 각 노드의 키 값을 자신의 수평 거리에 해당하는 맵(map) 버킷에 저장합니다.
  • C++의 map은 키를 기준으로 자동 정렬되므로, 왼쪽 열부터 오른쪽 열까지 순서대로 수직 라인을 처리할 수 있습니다.
  • 같은 열 안에서는 트리에서 위쪽에 있는 노드가 먼저 저장되므로, 전체 수직 순회 순서가 자연스럽게 유지됩니다.
  • 마지막으로 저장된 순서대로 노드를 카운트하여 K번째 노드를 찾아 반환하고, 끝까지 찾지 못하면 -1을 반환합니다.

C++ 구현 예제

#include<iostream>
#include<map>
#include<vector>
#include<queue>
using namespace std;
class Node {
    public:
    int key;
    Node *left, *right;
};
Node* getNode(int key){
    Node* node = new Node;
    node->key = key;
    node->left = node->right = NULL;
    return node;
}
int findKthNodeVertical(Node* root, int k) {
    if (!root || k == 0)
        return -1;
    int n = 0;
    int k_node = -1;
    map<int, vector<int> > current_map;
    int hd = 0;
    queue<pair<Node*, int> > que;
    que.push(make_pair(root, hd));
    while (!que.empty()) {
        pair<Node*, int> temp = que.front();
        que.pop();
        hd = temp.second;
        Node* node = temp.first;
        current_map[hd].push_back(node->key);
        if (node->left != NULL)
            que.push(make_pair(node->left, hd - 1));
        if (node->right != NULL)
            que.push(make_pair(node->right, hd + 1));
    }
    map<int, vector<int> >::iterator it;
    for (it = current_map.begin(); it != current_map.end(); it++) {
        for (int i = 0; i < it->second.size(); ++i) {
            n++;
            if (n == k)
                return (it->second[i]);
        }
    }
    if (k_node == -1)
        return -1;
}
int main() {
    Node* root = getNode(1);
    root->left = getNode(2);
    root->right = getNode(3);
    root->left->left = getNode(4);
    root->left->right = getNode(5);
    root->right->left = getNode(6);
    root->right->right = getNode(7);
    root->right->left->right = getNode(8);
    root->right->right->right = getNode(9);
    int k = 3;
    cout << "Kth node in vertical traversal: " << findKthNodeVertical(root, k);
}

실행 결과

Kth node in vertical traversal: 1

복잡도 분석

이 알고리즘은 트리의 모든 노드를 한 번씩 방문하므로, map 연산에 드는 로그 비용을 포함해 전체 시간 복잡도는 O(N log N)입니다. 여기서 N은 트리의 노드 수입니다. 공간 복잡도는 모든 노드를 맵과 큐에 저장해야 하므로 O(N)입니다.

마무리

이처럼 BFS와 맵을 조합하면 이진 트리의 수직 순회를 효율적으로 구현할 수 있고, 순회 도중 카운트만 추가하면 K번째 노드도 손쉽게 찾을 수 있습니다. 수평 거리 개념을 활용하면 위에서 아래로, 왼쪽에서 오른쪽으로 이어지는 수직 순회 순서를 명확하게 보장할 수 있다는 점이 이 접근 방식의 핵심입니다.