#include <iostream> #include <algorithm> #include <vector> #include <queue> using namespace std; struct Vertex { int a, b; vector<int> n; Vertex(): a(-1), b(-1) { } ~Vertex() { } }; typedef vector<Vertex> Graph; int main() { ios_base::sync_with_stdio(false); cin.tie(NULL); int n, m; cin >> n >> m; Graph G(n); for(int i = 0; i < n - 1; ++i) { int u, v; cin >> u >> v; --u; --v; G[u].n.push_back(v); G[v].n.push_back(u); } for(int u = 0; u < m; ++u) { int t; cin >> t; G[u].a = G[u].b = t; } queue<int> Q; for(int i = 0; i < m; ++i) Q.push(i); int last = 0; while(not Q.empty()) { int u = Q.front(); Q.pop(); last = u; if(G[u].a == -2) { vector<int> kek; for(unsigned int i = 0; i < G[u].n.size(); ++i) { int v = G[u].n[i]; if(G[v].a >= 0) { kek.push_back(G[v].a); kek.push_back(G[v].b); } } nth_element(kek.begin(), kek.begin() + kek.size() / 2 - 1, kek.end()); G[u].a = kek[kek.size() / 2 - 1]; nth_element(kek.begin(), kek.begin() + kek.size() / 2, kek.end()); G[u].b = kek[kek.size() / 2]; } for(unsigned int i = 0; i < G[u].n.size(); ++i) { int v = G[u].n[i]; if(G[v].a != -1) continue; G[v].a = -2; Q.push(v); } } vector<bool> queued(n, false); queue<int> R; R.push(last); queued[last] = true; int t = (G[last].b + G[last].a) / 2; G[last].b = G[last].a = t; int64_t result = 0; while(not R.empty()) { int u = R.front(); R.pop(); for(unsigned int i = 0; i < G[u].n.size(); ++i) { int v = G[u].n[i]; if(queued[v]) continue; int t = G[v].a; if(abs(G[v].a - G[u].a) > abs(G[v].b - G[u].a)) t = G[v].b; if(G[v].a <= G[u].a and G[u].a <= G[v].b) t = G[u].a; G[v].a = G[v].b = t; result += abs(G[v].a - G[u].a); R.push(v); queued[v] = true; } } cout << result << '\n'; return 0; }
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 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | #include <iostream> #include <algorithm> #include <vector> #include <queue> using namespace std; struct Vertex { int a, b; vector<int> n; Vertex(): a(-1), b(-1) { } ~Vertex() { } }; typedef vector<Vertex> Graph; int main() { ios_base::sync_with_stdio(false); cin.tie(NULL); int n, m; cin >> n >> m; Graph G(n); for(int i = 0; i < n - 1; ++i) { int u, v; cin >> u >> v; --u; --v; G[u].n.push_back(v); G[v].n.push_back(u); } for(int u = 0; u < m; ++u) { int t; cin >> t; G[u].a = G[u].b = t; } queue<int> Q; for(int i = 0; i < m; ++i) Q.push(i); int last = 0; while(not Q.empty()) { int u = Q.front(); Q.pop(); last = u; if(G[u].a == -2) { vector<int> kek; for(unsigned int i = 0; i < G[u].n.size(); ++i) { int v = G[u].n[i]; if(G[v].a >= 0) { kek.push_back(G[v].a); kek.push_back(G[v].b); } } nth_element(kek.begin(), kek.begin() + kek.size() / 2 - 1, kek.end()); G[u].a = kek[kek.size() / 2 - 1]; nth_element(kek.begin(), kek.begin() + kek.size() / 2, kek.end()); G[u].b = kek[kek.size() / 2]; } for(unsigned int i = 0; i < G[u].n.size(); ++i) { int v = G[u].n[i]; if(G[v].a != -1) continue; G[v].a = -2; Q.push(v); } } vector<bool> queued(n, false); queue<int> R; R.push(last); queued[last] = true; int t = (G[last].b + G[last].a) / 2; G[last].b = G[last].a = t; int64_t result = 0; while(not R.empty()) { int u = R.front(); R.pop(); for(unsigned int i = 0; i < G[u].n.size(); ++i) { int v = G[u].n[i]; if(queued[v]) continue; int t = G[v].a; if(abs(G[v].a - G[u].a) > abs(G[v].b - G[u].a)) t = G[v].b; if(G[v].a <= G[u].a and G[u].a <= G[v].b) t = G[u].a; G[v].a = G[v].b = t; result += abs(G[v].a - G[u].a); R.push(v); queued[v] = true; } } cout << result << '\n'; return 0; } |