[백준,c++] 20924번 - 트리의 기둥과 가지

문제

20924번: 트리의 기둥과 가지

 

코드

#include <iostream>
#include <vector>
#include <memory.h>

using namespace std;

vector<pair<int,int>>tree[200001];
bool visited[200001];
int trunk_len, start_branch, branch_len, end_branch;

void find_trunk(int node, int len){
    visited[node] = true;
    trunk_len = len;
    start_branch = node;
    if(tree[node].size()>2) return;
    
    for(auto next: tree[node]){
        if(!visited[next.first]){
            visited[next.first] = true;
            find_trunk(next.first,len+next.second);
        }
    }
}

void find_branch(int node, int len){
    visited[node] = true;
    if(len>branch_len){
        end_branch = node;
        branch_len = len;
    }
    
    for(auto next: tree[node]){
        if(!visited[next.first]){
            visited[next.first] = true;
            find_branch(next.first,len+next.second);
        }
    }
}

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(0);
    
    int N, R; cin>>N>>R;
    
    for(int i=0; i<N-1; i++){
        int a, b, c; cin>>a>>b>>c;
        tree[a].push_back({b,c});
        tree[b].push_back({a,c});
    }
    if(tree[R].size()>=2){
        find_trunk(R, 0);
        trunk_len = 0;
        memset(visited, 0, sizeof(visited));
        find_branch(R, 0);
    }
    else{
        find_trunk(R, 0);
        find_branch(start_branch, 0);
    }

    cout<<trunk_len<<' '<<branch_len;
     
}

 

풀이

문제의 조건 그대로 dfs를 이용하면 쉽게 해결할 수 있는 문제입니다.

1. 루트 노드는 문제에서 입력으로 주어지기 때문에 루트 노드부터 dfs를 출발해서 기둥까지 길이를 구해주고, 기둥이 끝나는 지점 또한 기록합니다.
2. 그다음 기둥이 끝나는 지점에서 가장 긴 리프 노드까지의 길이를 구하면 됩니다.

 

엣지 케이스

3 1
1 2 5
1 3 5
    
답 : 0 5

위와 같이 기둥이 없고, 가지고 2개인 경우를 잘 처리해주어야 합니다.