BOJ12912: 트리 수정

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

  • 트리의 지름 확장판 문제. 일단 시간 내에 들어오는 N 스퀘어 코드를 스스로 짜는데에 성공했다.
  • 일단 문제가 시키는데로 간선을 하나 지우면,
  • 두개의 서브트리가 나오는데, 두 트리의 지름을 구한 뒤, 지운 간선을 이용해 두 지름을 붙이면 최대인 것을 파악하는게 포인트다.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
#include <bits/stdc++.h>
using namespace std;

typedef long long LL;


struct Edge {
  LL v, c;
};

struct Pair {
  int u, ue, v, ve;
};

class Tree {
  int N;
  vector<vector<Edge>> adj;
  vector<Pair> edges;
public:
  Tree(int n): N(n), adj(n) {}
  void input();
  Edge farthest(LL u, LL p);
  LL solve();
};

void Tree::input() {
  for(int i = 0; i < N-1; i++) {
    int u, v, c;
    scanf("%d%d%d", &u, &v, &c);
    adj[u].push_back({v, c});
    adj[v].push_back({u, c});
    edges.push_back({u, adj[u].size()-1, v, adj[v].size()-1});
  }
}

Edge Tree::farthest(LL u, LL p) {
  Edge ans = {u, 0};
  for(auto [v,c] : adj[u]) {
		if(p == v || v == -1) continue;
    Edge cur = farthest(v, u);
    if(ans.c < c + cur.c)
      ans = {cur.v, cur.c + c};
  }
  return ans;
}

LL Tree::solve() {
  LL ans = 0;
  for(auto [u, ue, v, ve] : edges) {
    LL t1 = -1, t2 = -1;
    Edge &e1 = adj[u][ue];
    Edge &e2 = adj[v][ve];
    swap(t1, e1.v);
    swap(t2, e2.v);

    Edge point1 = farthest(t1, -1);
    LL radius1 = farthest(point1.v, -1).c;

    Edge point2 = farthest(t2, -1);
    LL radius2 = farthest(point2.v, -1).c;

    swap(t1, e1.v);
    swap(t2, e2.v);
    ans = max(ans, e1.c + radius1 + radius2);
  }
  return ans;
}

int main() {
  int N;
  scanf("%d", &N);
  Tree tree(N);
  tree.input();
  cout << tree.solve();
}