이진 트리와 특정 목표 노드(target), 그리고 값 K가 주어졌을 때, 목표 노드로부터 거리 K만큼 떨어져 있는 모든 노드의 값을 찾는 문제입니다.
예를 들어, 입력이 다음과 같다고 가정해 보겠습니다.
- root = [3,5,1,6,2,0,8,null,null,7,4]
- target = 5
- K = 2
이 경우 출력은 [7, 4, 1]이 됩니다. 목표 노드 5에서 거리 2만큼 떨어져 있는 노드들의 값이 각각 7, 4, 1이기 때문입니다.
문제 해결 접근 방식
이 문제의 핵심은 트리를 그래프처럼 다루는 것입니다. 일반적인 트리 순회는 부모에서 자식으로만 이동할 수 있지만, 거리 K를 계산하려면 자식에서 부모 방향으로도 이동할 수 있어야 합니다. 따라서 다음 두 단계로 문제를 해결합니다.
1단계: DFS로 각 노드의 부모 정보 저장
먼저 깊이 우선 탐색(DFS)을 수행하여 모든 노드의 부모 노드를 맵(map)에 기록합니다. 이렇게 하면 나중에 위쪽 방향으로도 탐색을 확장할 수 있습니다.
- dfs(node, pa) 함수를 정의합니다. pa는 부모 노드이며 기본값은 NULL입니다.
- 노드가 null이면 즉시 반환합니다.
- parent[node] := pa 로 현재 노드의 부모를 저장합니다.
- 왼쪽 자식과 오른쪽 자식에 대해 재귀적으로 dfs를 호출합니다.
2단계: BFS로 목표 노드에서 거리 K 탐색
부모 정보가 준비되면, 목표 노드를 시작점으로 너비 우선 탐색(BFS)을 수행합니다.
- 결과를 담을 배열 ans를 선언합니다.
- (노드, 거리) 쌍을 저장할 큐 q를 만들고 { target, 0 }을 삽입합니다.
- 방문 여부를 추적하기 위한 집합 visited를 만들고 target을 추가합니다.
- 큐가 빌 때까지 다음 과정을 반복합니다.
- 큐에서 요소를 꺼내 level(거리)과 node를 추출합니다.
- level == k 라면 해당 노드의 값을 ans에 추가합니다.
- 왼쪽 자식이 존재하고, 아직 방문하지 않았으며, level + 1 <= k 라면 큐에 { 왼쪽 자식, level + 1 }을 삽입하고 방문 처리합니다.
- 오른쪽 자식도 동일한 조건으로 처리합니다.
- 부모 노드가 존재하고, 아직 방문하지 않았으며, level + 1 <= k 라면 큐에 { 부모 노드, level + 1 }을 삽입하고 방문 처리합니다.
- 탐색이 끝나면 ans를 반환합니다.
BFS의 특성상 같은 거리(level)의 노드들이 순서대로 처리되므로, level이 K에 도달하는 순간의 노드들이 바로 정답이 됩니다. visited 집합을 사용하는 이유는 부모 방향과 자식 방향을 오가며 같은 노드를 중복 방문하는 것을 방지하기 위함입니다.
C++ 구현 예제
다음 코드를 통해 전체 구현을 확인해 보겠습니다.
#include <bits/stdc++.h>
using namespace std;
void print_vector(vector<int> v){
cout << "[";
for(int i = 0; i<v.size(); i++){
cout << v[i] << ", ";
}
cout << "]"<<endl;
}
class TreeNode{
public:
int val;
TreeNode *left, *right;
TreeNode(int data){
val = data;
left = NULL;
right = NULL;
}
};
void insert(TreeNode **root, int val){
queue<TreeNode*> q;
q.push(*root);
while(q.size()){
TreeNode *temp = q.front();
q.pop();
if(!temp->left){
if(val != NULL)
temp->left = new TreeNode(val);
else
temp->left = new TreeNode(0);
return;
}else{
q.push(temp->left);
}
if(!temp->right){
if(val != NULL)
temp->right = new TreeNode(val);
else
temp->right = new TreeNode(0);
return;
}else{
q.push(temp->right);
}
}
}
TreeNode *make_tree(vector<int> v){
TreeNode *root = new TreeNode(v[0]);
for(int i = 1; i<v.size(); i++){
insert(&root, v[i]);
}
return root;
}
class Solution {
public:
map <TreeNode*, TreeNode*> parent;
void dfs(TreeNode* node, TreeNode* pa = NULL){
if (!node)
return;
parent[node] = pa;
dfs(node->left, node);
dfs(node->right, node);
}
vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
vector<int> ans;
parent.clear();
dfs(root);
queue<pair<TreeNode*, int> > q;
q.push({ target, 0 });
set<TreeNode*> visited;
visited.insert(target);
while (!q.empty()) {
pair<TreeNode*, int> temp = q.front();
q.pop();
int level = temp.second;
TreeNode* node = temp.first;
if (level == k) {
ans.push_back(node->val);
}
if ((node->left && node->left->val != 0) && level + 1 <= k && !visited.count(node->left)) {
q.push({ node->left, level + 1 });
visited.insert(node->left);
}
if ((node->right && node->right->val != 0) && level + 1 <= k && !visited.count(node->right)){
q.push({ node->right, level + 1 });
visited.insert(node->right);
}
if (parent[node] != NULL && level + 1 <= k && !visited.count(parent[node])) {
q.push({ parent[node], level + 1 });
visited.insert(parent[node]);
}
}
return ans;
}
};
main(){
Solution ob;
vector<int> v = {3,5,1,6,2,0,8,NULL,NULL,7,4};
TreeNode *root = make_tree(v);
TreeNode *target = root->left;
print_vector(ob.distanceK(root, target, 2));
}입력
{3,5,1,6,2,0,8,NULL,NULL,7,4}출력
[7, 4, 1]
복잡도 분석
- 시간 복잡도: O(N) — DFS로 부모 정보를 저장하는 데 N, BFS로 전체 노드를 한 번씩 방문하므로 N입니다.
- 공간 복잡도: O(N) — 부모 맵, 방문 집합, 큐에 최대 N개의 노드가 저장될 수 있습니다.