C++로 트리의 특정 서브트리 DFS 순회에서 K번째 노드 찾는 방법

문제 개요

크기가 N인 트리와 트리의 한 노드 V, 그리고 정수 k가 주어졌을 때, 노드 V를 루트로 하는 서브트리의 DFS(깊이 우선 탐색) 순회에서 k번째 노드를 찾는 것이 이 글의 목표입니다.
다시 말해, 정점 V에서 시작하는 DFS 순회 과정에서 k번째로 방문되는 노드를 구해야 합니다.

예시로 이해하기

입력:

V = 2, k = 3

출력: 1

설명:

노드 2에서 시작하는 DFS 순회: {2, 6, 1}
3번째 노드는 1입니다.

해결 접근 방법

가장 직관적인 해결책은 노드 V에서 시작하는 DFS 순회 전체를 구한 뒤, 그 결과에서 k번째 값을 꺼내는 것입니다.

이를 효율적으로 처리하려면 트리 전체를 단 한 번 순회하면서 다음 세 가지 정보를 미리 기록해 두면 됩니다.

dfsTraversalVector: 루트에서 시작한 DFS 순회에서 노드가 방문된 순서대로 저장된 배열
startIdx[node]: 해당 노드의 서브트리 순회가 시작되는 인덱스
endIdx[node]: 해당 노드의 서브트리 순회가 끝나는 인덱스

핵심은 노드 V의 서브트리가 DFS 순회 배열에서 [startIdx[V], endIdx[V]]라는 연속된 구간에 나타난다는 점입니다. 따라서 k번째 노드는 단순히 dfsTraversalVector[startIdx[V] + k - 1]로 계산할 수 있습니다. 만약 이 인덱스가 endIdx[V]를 초과한다면 서브트리에 k개의 노드가 존재하지 않는 것이므로 -1을 반환합니다.

C++ 구현 예제

위 접근 방식을 실제로 구현한 프로그램입니다.

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

int n;
vector<int> tree[N];
int currentIdx;
vector<int> startIdx, endIdx;
vector<int> dfsTraversalVector;

void insertEdge(int u, int v){
    tree[u].push_back(v);
    tree[v].push_back(u);
}

void findDfsTraversal(int ch, int par){
    dfsTraversalVector[currentIdx] = ch;
    startIdx[ch] = currentIdx++;
    for (auto c : tree[ch]) {
        if (c != par)
            findDfsTraversal(c, ch);
    }
    endIdx[ch] = currentIdx - 1;
}

int findKNodeDfsV(int v, int k){
    k = k + (startIdx[v] - 1);
    if (k <= endIdx[v])
        return dfsTraversalVector[k];
    return -1;
}

int main(){
    n = 9;
    insertEdge(5, 8);
    insertEdge(5, 2);
    insertEdge(5, 10);
    insertEdge(5, 3);
    insertEdge(2, 6);
    insertEdge(2, 1);
    insertEdge(3, 9);
    insertEdge(9, 7);

    startIdx.resize(n);
    endIdx.resize(n);
    dfsTraversalVector.resize(n);

    findDfsTraversal(5, 0);

    int v = 2, k = 3;
    cout << k << "-th node in DFS traversal of node " << v << " is " << findKNodeDfsV(v, k);
    return 0;
}

실행 결과

3-th node in DFS traversal of node 2 is 1

복잡도 분석

시간 복잡도: O(N) — 트리 전체를 한 번만 순회합니다.
공간 복잡도: O(N) — 순회 배열과 시작·끝 인덱스 배열을 저장해야 합니다.

특히 (V, k) 형태의 쿼리가 여러 번 들어오는 경우에도 초기 순회를 한 번만 수행해 두면 각 노드의 순회 구간을 모두 알아낼 수 있으므로, 이후의 모든 쿼리를 O(1)에 처리할 수 있다는 점이 이 방법의 가장 큰 장점입니다.

마무리

DFS 순회에서 서브트리는 항상 연속된 구간을 이룬다는 성질(오일러 투어, Euler Tour 기법)을 활용하면 k번째 노드 조회를 아주 간단한 인덱스 계산으로 해결할 수 있습니다. 트리 관련 알고리즘 문제에서 매우 자주 등장하는 핵심 패턴이므로 반드시 익혀두시길 바랍니다.