C++로 구현하는 N-ary 트리에서 다음으로 큰 요소 찾기

N-ary 트리(n-ary tree)는 각 노드가 최대 n개의 자식 노드를 가질 수 있는 트리 구조입니다. 이번 글에서는 하나의 숫자 x가 주어졌을 때, N-ary 트리 전체를 탐색하여 x보다 크면서 가장 작은 값, 즉 다음으로 큰 요소(next greater element)를 찾는 방법을 알아보겠습니다.

이 문제는 트리를 순회하면서 조건에 맞는 후보 값을 계속 갱신해 나가는 방식으로 해결할 수 있습니다.

알고리즘

• N-ary 트리를 생성합니다.
• 결과값을 저장할 변수를 초기화합니다.
• 다음으로 큰 요소를 찾는 재귀 함수를 작성합니다.
  • 현재 노드가 NULL이면 그대로 반환합니다.
  • 현재 노드의 데이터가 찾고자 하는 값(x)보다 큰지 확인합니다.
  • 크다면, 결과가 아직 비어 있거나 현재 노드의 값이 기존 결과보다 작은지 검사합니다.
  • 조건을 만족하면 결과를 현재 노드로 갱신합니다.
  • 현재 노드의 모든 자식 노드를 순회하며 재귀 함수를 호출합니다.

핵심은 트리를 순회하는 동안 "주어진 수보다 크면서 기존 결과보다 작은" 값을 발견할 때마다 결과를 갱신하는 것입니다. 이렇게 하면 탐색이 끝난 시점에 항상 정확한 '다음으로 큰 요소'를 얻을 수 있습니다.

C++ 구현

위 알고리즘을 C++로 구현한 코드는 다음과 같습니다.

#include <bits/stdc++.h>
using namespace std;
struct Node {
   int data;
   vector<Node*> child;
};
Node* newNode(int data) {
   Node* newNode = new Node;
   newNode->data = data;
   return newNode;
}
void findNextGreaterElement(Node* root, int x, Node** result) {
   if (root == NULL) {
      return;
   }
   if (root->data > x) {
      if (!(*result) || (*result)->data > root->data) {
         *result = root;
      }
   }
   int childCount = root->child.size();
   for (int i = 0; i < childCount; i++) {
      findNextGreaterElement(root->child[i], x, result);
   }
   return;
}
int main() {
   Node* root = newNode(10);
   root->child.push_back(newNode(12));
   root->child.push_back(newNode(23));
   root->child.push_back(newNode(45));
   root->child[0]->child.push_back(newNode(40));
   root->child[1]->child.push_back(newNode(33));
   root->child[2]->child.push_back(newNode(12));
   Node* result = NULL;
   findNextGreaterElement(root, 20, &result);
   cout << result->data << endl;
   return 0;
}

실행 결과

위 코드를 실행하면 다음과 같은 결과가 출력됩니다.

23

예제 트리에서 20보다 큰 값들은 23, 45, 40, 33이며, 그중 가장 작은 값인 23이 출력되는 것을 확인할 수 있습니다.