C++로 완전 이진 트리의 각 노드에 next(다음 오른쪽) 포인터 채우기

문제 개요

완전 이진 트리(perfect binary tree)가 주어졌다고 가정해 봅시다. 이 트리의 각 노드는 (data, left, right, next) 네 개의 필드를 가지고 있습니다. left는 왼쪽 서브트리를, right는 오른쪽 서브트리를 가리키며, next 포인터는 같은 레벨(level)에 있는 바로 오른쪽 이웃 노드를 가리켜야 합니다. 만약 오른쪽에 이웃 노드가 존재하지 않는다면 next는 NULL이 됩니다.

처음에는 모든 next 포인터가 NULL로 초기화되어 있으며, 우리의 목표는 이 링크들을 올바르게 연결하는 것입니다. 예를 들어 아래와 같은 트리가 있다면, 변환 후 각 노드는 자신의 오른쪽 이웃을 가리키게 됩니다.

해결 알고리즘

이 문제는 O(1) 추가 공간으로 레벨 순회(level-order traversal)를 수행하여 해결할 수 있습니다. 핵심 아이디어는 이미 연결된 상위 레벨의 next 포인터를 활용해 다음 레벨의 노드들을 차례대로 잇는 것입니다. 단계별 절차는 다음과 같습니다.

pre := root, nextPre := null, prev := null로 초기화합니다.
pre가 NULL이 아닌 동안 반복합니다.
- pre가 NULL이 아닌 동안 내부 반복을 수행합니다.
  - pre의 왼쪽 자식이 존재하면:
    - prev가 NULL이 아니라면 prev->next := pre->left로 설정하고, 그렇지 않으면 nextPre := pre->left로 설정합니다.
    - prev := pre->left로 갱신합니다.
  - pre의 오른쪽 자식이 존재하면:
    - prev가 NULL이 아니라면 prev->next := pre->right로 설정하고, 그렇지 않으면 nextPre := pre->right로 설정합니다.
    - prev := pre->right로 갱신합니다.
  - pre := pre->next로 현재 레벨의 다음 노드로 이동합니다.
- 현재 레벨의 순회가 끝나면 pre := nextPre로 다음 레벨의 시작 노드로 이동합니다.
- nextPre와 prev를 NULL로 초기화합니다.
모든 연결이 완료되면 루트를 반환합니다.

C++ 구현 예제

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

#include <bits/stdc++.h>
#include <stack>
using namespace std;
class Node {
public:
   int val;
   Node* left;
   Node* right;
   Node* next;
   Node() {}
   Node(int _val, Node* _left, Node* _right) {
      val = _val;
      left = _left;
      right = _right;
      next = NULL;
   }
};
class Solution {
public:
   Node* connect(Node* root) {
      Node* pre = root;
      Node* nextPre = NULL;
      Node* prev = NULL;
      while(pre){
         while(pre){
            //cout << pre->val << endl;
            if(pre->left){
               if(prev){
                  prev->next = pre->left;
               }else{
                  nextPre = pre->left;
               }
               prev = pre->left;
            }
            if(pre->right){
               if(prev){
                  prev->next = pre->right;
               }else{
                  nextPre = pre->right;
               }
               prev = pre->right;
            }
            pre = pre->next;
         }
         //cout << "*" << endl;
         pre = nextPre;
         nextPre = NULL;
         prev = NULL;
    }
      return root;
   }
};
void printTree(Node* root) {
   cout << "[";
   if (root == NULL) return;
   queue<Node*> q;
   Node *curr;
   q.push(root);
   q.push(NULL);
   while (q.size() > 1) {
      curr = q.front();
      q.pop();
      if (curr == NULL){
         q.push(NULL);
      }
      else {
         // if(curr->next)
         // q.push(curr->next);
         if(curr->left)
            q.push(curr->left);
            if(curr->right)
               q.push(curr->right);
               if(curr->val == 0){
                  cout << "null" << ", ";
               }else{
                  cout << curr->val << ", ";
                  if (curr->next == NULL) cout<<"#, ";
               }
      }
   }
   cout << "]"<<endl;
}
int main() {
Node* root;
Node nodeFour(4, NULL, NULL);
Node nodeFive(5, NULL, NULL );
Node nodeSeven(7, NULL, NULL);
Node nodeSix(6, NULL, NULL);
Node nodeTwo(2,&nodeFour,&nodeFive);
Node nodeThree(3,&nodeSix,&nodeSeven);
Node nodeOne(1,&nodeTwo,&nodeThree);
root = &nodeOne;
Solution ob;
root = ob.connect(root);
printTree(root);
}

입력

[1,2,3,4,5,6,7]
Node* root;
Node nodeFour(4, NULL, NULL);
Node nodeFive(5, NULL, NULL );
Node nodeSeven(7, NULL, NULL);
Node nodeSix(6, NULL, NULL);
Node nodeTwo(2,&nodeFour,&nodeFive);
Node nodeThree(3,&nodeSix,&nodeSeven);
Node nodeOne(1,&nodeTwo,&nodeThree);
root = &nodeOne;
Solution ob;
root = ob.connect(root);

출력

[1, #, 2, 3, #, 4, 5, 6, 7, #]

복잡도 분석

이 알고리즘은 트리의 모든 노드를 정확히 한 번씩 방문하므로 시간 복잡도는 O(N)입니다. 또한 큐나 재귀 없이 기존의 next 포인터만 활용하기 때문에 공간 복잡도는 O(1)로, 추가 메모리 사용 없이 문제를 해결할 수 있다는 장점이 있습니다.