[BOJ] 12912 - 트리 수정

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

BOJ - 트리 수정

문제 유형: 트리, 그래프, 트리의 지름

문제 난이도: Gold I

시간 제한: 2초

메모리 제한: 512MB

 

문제

N개의 정점으로 이루어진 트리 T가 있다. 트리의 각 정점은 0번부터 N-1번까지 번호가 매겨져 있다.

• 트리에서 임의의 두 정점을 연결하는 단순 경로의 개수는 1개이다.
• 두 정점사이의 거리는 두 정점을 연결하는 단순 경로상에 있는 간선의 가중치의 합이다.
• 트리의 지름은 트리에 존재하는 모든 경로 중에서 가장 긴 것이다.

홍준이는 T에서 간선을 하나 제거하고, 간선을 하나 추가하려고 한다. 이때, 추가하는 간선의 가중치는 제거한 간선의 가중치와 같아야 하며, 간선을 추가한 이후에도 트리를 유지해야 한다.

이때, 홍준이가 만들 수 있는 트리 중에서 지름이 가장 큰 것을 구하는 프로그램을 작성하시오.

 

입력

첫째 줄에 트리 정점의 개수 N이 주어진다. (2 ≤ N ≤ 2,000)

둘째 줄부터 N-1개의 줄에는 트리를 이루는 간선이 주어진다. 간선은 from, to, cost와 같이 세 가지 정수로 이루어져 있으며, from과 to를 연결하는 간선의 가중치가 cost라는 뜻이다. (1 ≤ cost ≤ 1,000,000,000)

 

출력

첫째 줄에 홍준이가 만들 수 있는 트리 중에서 가장 지름이 큰 것의 지름을 출력한다.

 

풀이

모든 간선을 한 번씩 제거해보는 경우를 따져보자.

간선을 하나 제거한다면, 두 개의 트리로 분리된다. 

제거한 간선을 기반으로 각각의 트리에 접근할 수 있게 된다.

두 트리의 각각의 지름을 구하고, 없앤 간선의 지름을 더한다. 이를 모든 간선을 한 번씩 독립적으로 지우는것에 대해 고려해서 구하면 된다.

 

트리에서의 각 노드간 경로는 유일하므로, 간선을 제거하는 것을 그 간선을 부속하는 노드 u, v에서 u에서 탐색하며 v를 막는 방법을 쓰면, u쪽으로 나눠진 부분의 분할된 트리를 탐색할 수 있다.

반대도 그렇다.

 

코드

C++

#include <bits/stdc++.h>

using namespace std;
using ll = long long;
using pll = pair<ll, ll>;

int n;

ll getDiameter(ll start, ll ban, vector<vector<pll>> &adj) {
    if (adj[start].size() == 1) return 0;

    vector<ll> dist(n, LLONG_MAX);
    dist[start] = 0;
    queue<ll> q;
    q.emplace(start);
    ll maxDist1 = 0;
    ll point1 = start;
    while (!q.empty()) {
        ll curr = q.front();
        q.pop();
        if (dist[curr] > maxDist1) {
            point1 = curr;
            maxDist1 = dist[curr];
        }
        for (const auto &[nextNode, nextCost] : adj[curr]) {
            if (nextNode == ban) continue;
            if (dist[nextNode] > dist[curr] + nextCost) {
                dist[nextNode] = dist[curr] + nextCost;
                q.emplace(nextNode);
            }
        }
    }
    vector<ll> dist2(n, LLONG_MAX);
    dist2[point1] = 0;
    q.emplace(point1);
    ll diameter = 0;
    ll point2 = point1;

    while (!q.empty()) {
        ll curr = q.front();
        q.pop();
        if (dist2[curr] > diameter) {
            point2 = curr;
            diameter = dist2[curr];
        }
        for (const auto &[nextNode, nextCost] : adj[curr]) {
            if (nextNode == ban) continue;
            if (dist2[nextNode] > dist2[curr] + nextCost) {
                dist2[nextNode] = dist2[curr] + nextCost;
                q.emplace(nextNode);
            }
        }
    }

    return diameter;
}

int main(int argc, char *argv[]) {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> n;
    vector<vector<pll>> adj(n);
    vector<array<ll, 3>> edges(n - 1);
    for (int i = 0; i < n - 1; i++) {
        ll u, v, w;
        cin >> u >> v >> w;
        edges[i] = {u, v, w};

        adj[u].emplace_back(v, w);
        adj[v].emplace_back(u, w);
    }

    ll ans = 0;
    for (int i = 0; i < n - 1; i++) {
        ll ban1 = edges[i][0];
        ll ban2 = edges[i][1];
        ll weight = edges[i][2];

        ll d1 = getDiameter(ban1, ban2, adj);
        ll d2 = getDiameter(ban2, ban1, adj);
        ans = max(ans, d1 + d2 + weight);
    }

    cout << ans << "\n";

    return 0;
}