C++에서 이진 트리의 잎 찾기: DFS 높이 계산으로 레벨별 노드 제거하기

문제 이해

하나의 이진 트리가 주어졌다고 가정해 봅시다. 우리는 트리의 모든 잎(leaf) 노드를 수집한 뒤 제거하고, 이 과정을 트리가 완전히 비워질 때까지 반복해야 합니다. 각 라운드마다 제거된 노드들을 차례대로 기록하면, 마치 트리를 아래층부터 한 겹씩 벗겨 내는 것과 같은 결과를 얻게 됩니다.

예를 들어 아래와 같은 이진 트리가 입력으로 주어지면,

출력은 다음과 같습니다.

[[4, 5, 3], [2], [1]]

첫 번째 라운드에서 잎 노드인 4, 5, 3이 제거되고, 두 번째 라운드에서 2가, 마지막 라운드에서 루트인 1이 제거됩니다. 참고로 같은 라운드 안에서 노드의 나열 순서는 탐색 방향에 따라 달라질 수 있습니다.

접근 방법: 노드의 높이(height) 활용하기

이 문제의 핵심은 "노드가 몇 번째 라운드에 제거되는가"와 "노드의 높이"가 서로 같다는 사실입니다. 여기서 높이란 해당 노드부터 그 아래에 있는 가장 깊은 잎 노드까지의 거리를 의미합니다.

• 높이가 1인 노드 → 첫 번째 라운드에서 제거되는 잎 노드
• 높이가 2인 노드 → 두 번째 라운드에서 제거
• 높이가 k인 노드 → k번째 라운드에서 제거

따라서 DFS(깊이 우선 탐색)로 모든 노드의 높이를 계산한 뒤, 높이가 같은 노드끼리 같은 그룹에 모아 주면 곧바로 정답을 구할 수 있습니다. 실제로 트리에서 노드를 삭제하는 과정은 필요하지 않습니다.

구체적인 알고리즘은 다음과 같습니다.

• 노드의 값을 키로, 높이를 값으로 저장할 맵 sz를 하나 정의합니다.
• 결과를 담을 2차원 배열 ret를 정의합니다.
• 노드를 인자로 받는 dfs() 함수를 정의합니다.
  • 노드가 NULL이면 0을 반환합니다.
  • sz[노드의 값] := 1 + max(dfs(왼쪽 자식), dfs(오른쪽 자식))로 높이를 계산합니다.
  • ret의 크기가 sz[노드의 값]보다 작으면 새로운 배열 temp를 만들어 ret의 끝에 추가합니다.
  • 노드의 값을 ret[sz[노드의 값] − 1]의 끝에 삽입합니다.
  • sz[노드의 값]을 반환합니다.
• 메인 함수에서 dfs(root)를 호출한 뒤 ret를 반환합니다.

C++ 구현 예제

아래 구현 코드를 통해 더 잘 이해해 보겠습니다.

#include <bits/stdc++.h>
using namespace std;

void print_vector(vector<vector<int>> v){
   cout << "[";
   for(int i = 0; i < v.size(); i++){
      cout << "[";
      for(int j = 0; j < v[i].size(); j++){
         cout << v[i][j] << ", ";
      }
      cout << "],";
   }
   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:
   unordered_map<int, int> sz;
   vector<vector<int>> ret;
   int dfs(TreeNode* node){
      if(!node) return 0;
      sz[node->val] = 1 + max(dfs(node->left), dfs(node->right));
      if(ret.size() < sz[node->val]){
         vector<int> temp;
         ret.push_back(temp);
      }
      ret[sz[node->val] - 1].push_back(node->val);
      return sz[node->val];
   }
   vector<vector<int>> findLeaves(TreeNode* root) {
      dfs(root);
      return ret;
   }
};

int main(){
   Solution ob;
   vector<int> v = {1,2,3,4,5};
   TreeNode *root = make_tree(v);
   print_vector(ob.findLeaves(root));
}

실행 결과

입력:

{1,2,3,4,5}

출력:

[[4, 5, 3],[2],[1]]

복잡도 분석 및 참고 사항

• 시간 복잡도: O(n) — 각 노드를 정확히 한 번씩만 방문합니다.
• 공간 복잡도: O(n) — 재귀 호출 스택과 결과 배열에 노드 개수에 비례하는 메모리가 사용됩니다.

한 가지 유의할 점은 위 코드에서 높이를 노드의 값(value)을 키로 하는 해시 맵에 저장한다는 것입니다. 트리에 중복된 값이 존재하면 키 충돌이 발생할 수 있으므로, 실전에서는 맵 대신 각 노드 객체에 높이를 직접 저장하거나 재귀 반환값만 활용하는 방식이 더 안전합니다.