백준 20924 트리의 기둥과 기지 (C++)

https://www.acmicpc.net/problem/20924

#include <iostream>
#include <vector>
#include <algorithm>

using namespace std;

//g[i][n]={dist, j}
vector<pair<int, int>> g[200001];

bool visit[200001];


int maxTrunk = 0;
int maxBranch = 0;

int giga;

void dfs_trunk(int cur, int len){

    visit[cur] = true;
    maxTrunk = max(maxTrunk, len);
    giga = cur;

    if (g[cur].size() > 2) {
        return;
    }

    for(int i=0;i<g[cur].size();i++){
        if(!visit[g[cur][i].second]) {
            visit[g[cur][i].second] = true;
            dfs_trunk(g[cur][i].second, len + g[cur][i].first);
        }
    }

}

void dfs_branch(int cur, int len){
    visit[cur] = true;
    maxBranch = max(maxBranch, len);

    for(int i=0;i<g[cur].size();i++){
        if(!visit[g[cur][i].second]) {
            visit[g[cur][i].second] = true;
            dfs_branch(g[cur][i].second, len + g[cur][i].first);
        }
    }
}


int main() {

    ios::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);

    int n, r;

    cin >> n >> r;


    for (int i = 0; i < n - 1; i++) {
        int a, b, d;
        cin >> a >> b >> d;

        g[a].push_back(make_pair(d, b));
        g[b].push_back(make_pair(d, a));
    }


    if (g[r].size() >= 2) {
        maxTrunk = 0;
        dfs_branch(r, 0);

    } else {
        dfs_trunk(r, 0);
        dfs_branch(giga, 0);
    }


    cout << maxTrunk << " " << maxBranch;

    return 0;
}

1. 문제

트리의 노드의 개수와 루트 노드, 각 edge와 길이가 주어진다. 

기둥과 가지를 위와 같이 정의할 때, 기둥의 길이와 가장 긴 가지의 길이를 구하라.

2. 풀이

dfs를 이용하여 풀 수 있다. 루트 노드가 주어졌으니 여기서 시작, 기가노드 까지의 거리를 구하고, 그 후 기가노드에서부터 각 leaf 까지의 거리 중 최댓값을 구한다.

    for (int i = 0; i < n - 1; i++) {
        int a, b, d;
        cin >> a >> b >> d;

        g[a].push_back(make_pair(d, b));
        g[b].push_back(make_pair(d, a));
    }

edge를 입력받아 그래프 벡터에 저장해준다.

void dfs_trunk(int cur, int len){

    visit[cur] = true;
    maxTrunk = max(maxTrunk, len);
    giga = cur;

    if (g[cur].size() > 2) {
        return;
    }

    for(int i=0;i<g[cur].size();i++){
        if(!visit[g[cur][i].second]) {
            visit[g[cur][i].second] = true;
            dfs_trunk(g[cur][i].second, len + g[cur][i].first);
        }
    }

}

기둥의 길이를 구하는 로직이다. 기가 노드는 루트 부터 순회할 때, 자식 노드가 2개 이상인 첫 노드이다. 순회하면서, 조건에 맞다면 로직을 종료한다.

void dfs_branch(int cur, int len){
    visit[cur] = true;
    maxBranch = max(maxBranch, len);

    for(int i=0;i<g[cur].size();i++){
        if(!visit[g[cur][i].second]) {
            visit[g[cur][i].second] = true;
            dfs_branch(g[cur][i].second, len + g[cur][i].first);
        }
    }
}

가지의 최대 길이를 구하는 로직이다. 각 자식 노드를 돌면서 방문하지 않았을 경우 가지의 길이를 갱신하여 방문한다.

    if (g[r].size() >= 2) {
        maxTrunk = 0;
        dfs_branch(r, 0);

    } else {
        dfs_trunk(r, 0);
        dfs_branch(giga, 0);
    }


    cout << maxTrunk << " " << maxBranch;

 다시 메인함수이다. 기가 노드와 루트 노드가 같은 경우를 처리해주어야 한다. 만약 같다면 루트 노드의 자식이 둘 이상이고, 이 때 기둥의 길이는 0이다. 또, 이는 기가노드이기도 하므로 여기서부터 가지의 길이를 구해준다. 

 그렇지 않은 경우 루트 노드에서 기가 노드까지의  거리를 구하고, 기가 노드로부터 각 leaf까지의 거리의 최댓값을 구해준다.

 연산 후 해당 길이들을 출력하면 정답을 구할 수 있다.