Computer >> 컴퓨터 >  >> 프로그래밍 >> C++

트리 서브트리 DFS 쿼리를 C++로 효율적으로 처리하는 방법

이 문제에서는 하나의 트리가 주어지며, 지정된 특정 노드를 루트로 간주하여 해당 노드로부터 DFS(깊이 우선 탐색)를 수행해야 합니다.

문제 상황

예를 들어 아래 트리에서 노드 F를 기준으로 DFS를 수행해야 한다고 가정해 보겠습니다. 일반적인 방식이라면 쿼리가 들어올 때마다 해당 노드에서 DFS를 매번 새로 실행해야 하지만, 이는 비효율적입니다.

효율적인 접근 방식

이 튜토리얼에서는 시간 복잡도를 크게 줄일 수 있는 비전통적인 기법을 적용하여, 제약 조건이 큰 입력에서도 코드가 시간 초과(TLE) 없이 동작하도록 만들어 보겠습니다.

모든 쿼리마다 DFS를 단순히 수행하는 나이브한 방식은 높은 제약 조건에서 실패하기 때문에, 우리는 DFS 방문 순서를 한 번만 미리 계산해 두는 방법을 사용합니다.

핵심 아이디어

  • 전체 트리에 대해 DFS를 한 번 수행하면서 방문 순서를 배열에 저장합니다.
  • 동시에 각 노드를 루트로 하는 서브트리에 포함된 노드의 개수(nodesunder)를 계산합니다.
  • 쿼리가 들어오면 해당 노드가 DFS 배열에서 차지하는 시작 인덱스부터 서브트리 크기만큼의 구간을 출력합니다.
#include <bits/stdc++.h>
using namespace std;
#define N 100000

// 트리 노드의 연결 관계를 저장하는 인접 리스트
vector<int> v[N];
unordered_map<int, int> mape; // 노드와 DFS 배열상의 인덱스를 연결하는 맵
vector<int> a;

// DFS를 수행하면서 각 노드의 서브트리 크기(nodesunder)를 사전 계산하는 함수
void dfs(int nodesunder[], int child, int parent){
    a.push_back(child); // 트리의 DFS 순서를 저장
    nodesunder[child] = 1; // 자식 노드의 서브트리 크기 초기화
    for (auto it : v[child]) { // 일반적인 DFS 수행
        if (it != parent) {
            // 자식이 부모로 거슬러 올라가면 사이클이 생기므로 이를 방지
            dfs(nodesunder, it, child); // 재귀 호출
            nodesunder[child] += nodesunder[it]; // 자식들의 서브트리 크기를 누적
        }
    }
}

// 특정 노드의 서브트리에 대한 DFS를 출력하는 함수
void printDFS(int node, int nodesunder[]){
    int ind = mape[node]; // DFS 배열에서 해당 노드의 인덱스
    cout << "The DFS of subtree " << node << ": ";
    // DFS 배열을 순회하며 주어진 노드 아래의 모든 노드를 출력
    for (int i = ind; i < ind + nodesunder[node]; i++){
        cout << a[i] << " ";
    }
    cout << endl;
}

// 인접 리스트를 유지하기 위한 함수
void addEdgetoGraph(int x, int y){
    v[x].push_back(y);
    v[y].push_back(x);
}

// DFS 배열에서 각 노드의 인덱스를 맵에 기록하는 함수
void mark(){
    int size = a.size();
    for (int i = 0; i < size; i++) {
        mape[a[i]] = i;
    }
}

int main(){
    int n = 7;
    // 트리의 간선 추가
    addEdgetoGraph(1, 2);
    addEdgetoGraph(1, 3);
    addEdgetoGraph(2, 4);
    addEdgetoGraph(2, 5);
    addEdgetoGraph(4, 6);
    addEdgetoGraph(4, 7);

    // 모든 노드의 서브트리에 포함된 노드 수를 저장하는 배열
    int nodesunder[n + 1];
    dfs(nodesunder, 1, 0); // nodesunder 배열 생성
    mark(); // 맵에 인덱스 기록

    // 쿼리 1
    printDFS(2, nodesunder);
    // 쿼리 2
    printDFS(4, nodesunder);
    return 0;
}

실행 결과

The DFS of subtree 2: 2 4 6 7 5
The DFS of subtree 4: 4 6 7

코드 이해하기

이 접근 방식의 핵심은 DFS 방문 순서를 미리 계산하여 벡터에 저장해 두는 것입니다. DFS를 사전 계산하는 과정에서 각 노드를 루트로 하는 서브트리에 포함된 노드의 개수(nodesunder)도 함께 구해 둡니다.

이후 쿼리가 들어오면 매번 새로 DFS를 수행할 필요 없이, 해당 노드의 DFS 배열상 시작 인덱스부터 서브트리에 포함된 노드 수만큼의 구간만 순회하여 출력하면 됩니다. 덕분에 사전 계산 이후에는 각 쿼리가 매우 빠르게 처리되며, 쿼리 개수가 많아져도 효율적으로 대응할 수 있습니다.

마무리

이 튜토리얼에서는 트리에서 특정 노드의 서브트리에 대한 DFS를 구하는 쿼리 문제를 다루었습니다. 사전 계산(precomputation) 기반의 접근 방식을 통해 반복적인 DFS 호출로 인한 시간 낭비를 줄이는 방법을 배웠습니다.

동일한 로직은 C, Java, Python 등 다른 프로그래밍 언어로도 손쉽게 구현할 수 있습니다. 이 글이 여러분의 학습에 도움이 되었기를 바랍니다.