백준 20924번 트리의 기둥과 가지

문제 요약

트리와 루트노드가 주어진다. 처음으로 자식의 수가 2 이상이 되는 노드를 기가 노드라고 하자. 만약에 그런 노드가 없다면 리프 노드를 기가 노드라고 하자.

문제에서 원하는 것은 루트에서 기가 노드까지의 거리와 기가 노드를 루트로 하는 서브트리에서 기가 노드와의 거리가 가장 먼 노드의 거리를 출력해야 한다.

노드 수 N은 최대 20만이다.

풀이

먼저 기가 노드를 찾아야 한다. 이는 루트 노드에서 부터 dfs를 진행하면서 자식의 수가 2 이상인 노드를 찾으면 dfs를 종료하는 식으로 찾을 수 있다.

기가 노드에서부터 가장 먼 노드를 찾는 것도 위 과정에서 기가 노드의 부모만 잘 저장해뒀다가 dfs를 돌리면 찾을 수 있다.

#include<bits/stdc++.h>
using namespace std;
const int MAXN = 2e5+5;
int N, R;
vector<vector<pair<int,int>>> G(MAXN);
int children[MAXN];
int depth[MAXN];
int giga_node, giga_parent;
void get_depth(int cur, int par, int d) {
    depth[cur] = d;
    for(pair<int,int> p : G[cur]) {
        int v = p.first;
        int w = p.second;
        if(v == par) continue;
        get_depth(v, cur, d+w);
    }
    return ;
}


void get_children(int cur, int par) {
    for(pair<int,int> p : G[cur]) {
        int v = p.first;
        if(v == par) continue;
        children[cur]++;
    }
    if(children[cur] > 1) {
        giga_node = cur;
        giga_parent = par;
        return ;
    }
    for(pair<int,int> p : G[cur]) {
        int v = p.first;
        if(v == par) continue;
        if(giga_node) return ;
        get_children(v, cur);
    }
    return ;
}

int dfs(int cur, int par) {
    int ret = 0;
    for(pair<int,int> p : G[cur]) {
        int v = p.first;
        int w = p.second;
        if(v == par) continue;
        ret = max(ret, dfs(v, cur) + w);
    }
    return ret;
}

int main() {
    cin.tie(nullptr); ios::sync_with_stdio(false);
    cin >> N >> R;
    for(int i=1;i<N;++i) {
        int u, v, w; cin >> u >> v >> w;
        G[u].emplace_back(v, w);
        G[v].emplace_back(u, w);
    }
    get_depth(R, 0, 0);
    get_children(R, 0);
    if(giga_node == 0) {
        for(int i=1;i<=N;++i) if(children[i] == 0) giga_node = i;
        cout << depth[giga_node] << ' ' << 0 << '\n';
        return 0;
    }
    cout << depth[giga_node] << ' ' << dfs(giga_node, giga_parent) << '\n';
    return 0;   
}