[코드트리] 트리 최적의 노드 C++

문제

https://www.codetree.ai/missions/9/problems/tree-optimal-node?&utm_source=clipboard&utm_medium=text

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

 

코드

#include <iostream>
#include <vector>
#include <cstring> //memset()

using namespace std;

const int MAX_N = 100001;

int n, max_dist, max_node;
vector<int> edges[MAX_N];
bool visited[MAX_N];

void DFS(int x, int d){
    if(d>max_dist){
        max_dist=d;
        max_node=x;
    }

    for(int i=0; i<(int)edges[x].size(); ++i){
        int y=edges[x][i];
        if(visited[y]) continue;
        visited[y]=true;
        DFS(y, d+1);
    }
}

int main() {
    // input
    cin >> n;
    for(int i=1; i<n; ++i){
        int a, b;
        cin >> a >> b;
        // a, b 노드 연결(모든 간선의 길이는 1)
        edges[a].push_back(b);
        edges[b].push_back(a);
    }

    // solution
    // 가장 먼 노드(a)에서 가장 먼 노드(b)의 길이를 구하면
    // a, b의 중간에 있는 노드가 가장 먼 노드까지의 거리가 최소인 노드
    DFS(1, 0);

    max_dist=0;
    memset(visited, false, sizeof(visited));
    DFS(max_node, 0);

    // output
    // 정중앙에 있는 노드에서 가장 먼 노드까지의 길이
    cout << (max_dist+1)/2 << '\n';

    return 0;
}

 

(+) 시간초과 코드

처음에는 한 지점에서 시작해서 가장 먼 노드의 길이를 다 구한 뒤, 그 중 최소값을 출력하려고 생각했다.

갑자기 트리의 지름을 구해서 그 중간에 있는 노드의 거리를 구하면 될 것 같다는 생각이 들었다.

#include <iostream>
#include <vector>
#include <cstring> //memset()

using namespace std;

const int MAX_N = 100001;

int n, max_dist;
vector<int> edges[MAX_N];
bool visited[MAX_N];
int dist[MAX_N]; // dist[i]: i번 노드에서 출발했을 때 가장 먼 노드까지의 거리

// s에서 x노드까지의 거리 d
void DFS(int s, int x, int d){
    
    if(d>max_dist){
        max_dist=d;
        dist[s]=max_dist;
    }

    for(int i=0; i<(int)edges[x].size(); ++i){
        int y=edges[x][i];
        if(visited[y]) continue;
        visited[y]=true;
        DFS(s, y, d+1);
    }
}

int main() {
    // input
    cin >> n;
    for(int i=1; i<n; ++i){
        int a, b;
        cin >> a >> b;
        // a, b 노드 연결(모든 간선의 길이는 1)
        edges[a].push_back(b);
        edges[b].push_back(a);
    }

    // solution
    for(int i=1; i<=n; ++i){
        max_dist=0;
        memset(visited, false, sizeof(visited));
        // i에서 i까지 거리가 0임->이후 i에 연결된 노드를 탐색
        visited[i]=true;
        DFS(i, i, 0);
    }

    // output
    int answer=1e9;
    for(int i=1; i<=n; ++i){
        answer=min(answer, dist[i]);
    }
    cout << answer << '\n';
    return 0;
}