#include <iostream>
#include <vector>
#include <algorithm>
using namespace std;
const int N = 500001;
int H[N], W[N], LV[N], mi[N], ma[N];
vector<int> G[N], L[N];
int n, m;
bool vis[N];
int main()
{
ios_base::sync_with_stdio(0);
cin >> n >> m;
for (int i = 0; i < n - 1; i++) {
int a, b;
cin >> a >> b;
G[a].push_back(b);
G[b].push_back(a);
LV[i + 2] = 2 * n;
}
for (int i = 1; i <= m; i++) {
cin >> W[i];
L[0].push_back(i);
LV[i] = 0;
mi[i] = W[i];
ma[i] = W[i];
}
int cnt = m;
int k = 0;
while (cnt != n) {
for (int i = 0; i < L[k].size(); i++) {
int leaf = L[k][i];
for (int j = 0; j < G[leaf].size(); j++) {
int parent = G[leaf][j];
if (LV[parent] != LV[leaf]) {
H[parent]++;
if (H[parent] == G[parent].size() - 1) {
L[k + 1].push_back(parent);
cnt++;
LV[parent] = k + 1;
}
}
}
}
k++;
}
// cout << "LEVELS" << endl;
for (int i = 1; i <= n; i++) {
// cout << LV[i] << endl;
}
long long res = 0;
for (int i = 1; i <= k; i++) {
for (int j = 0; j < L[i].size(); j++) {
int parent = L[i][j];
// cout << "ANALYZING " << parent << endl;
vector<int> children_min;
vector<int> children_max;
vector<long long> children_all;
for (int l = 0; l < G[parent].size(); l++) {
int leaf = G[parent][l];
if (LV[leaf] >= i) {
continue;
}
children_min.push_back(mi[leaf]);
children_max.push_back(ma[leaf]);
children_all.push_back(mi[leaf]);
children_all.push_back(ma[leaf]);
// cout << "ADDING " << mi[leaf] << " " << ma[leaf] << endl;
}
sort(children_min.begin(), children_min.end());
sort(children_max.begin(), children_max.end());
sort(children_all.begin(), children_all.end());
long long mini = 0, maxi = 0, result = 0, best = 999999999999999999, bestmin, bestmax;
long long a = 0, b = 0;
for (a = 0; a < children_min.size(); a++) {
result += children_min[a] - children_all[0];
}
// cout << "STARTING RESULT " << result << endl;
a = 0;
for (int l = 1; l < children_all.size(); l++) {
// cout << "POINT NUMBER " << l << endl;
while (a < children_min.size() && children_min[a] < children_all[l]) {
a++;
// cout << "A UP " << a << endl;
}
while (b < children_max.size() && children_max[b] < children_all[l]) {
b++;
// cout << "B UP " << b << endl;
}
result += b * (children_all[l] - children_all[l - 1]);
// cout << "GOING UP BY " << b * (children_all[l] - children_all[l - 1]) << endl;
result -= (children_min.size() - a) * (children_all[l] - children_all[l - 1]);
// cout << "GOING DOWN BY " << (children_min.size() - a) * (children_all[l] - children_all[l - 1]) << endl;
// cout << "CURRENT RESULT " << result << endl;
if (result < best) {
best = result;
bestmin = children_all[l];
bestmax = children_all[l];
} else if (result == best) {
bestmax = children_all[l];
}
}
mi[parent] = bestmin;
ma[parent] = bestmax;
res += best;
// cout << "FINAL RESULT " << res << endl;
// cout << mi[parent] << " " << ma[parent] << endl << endl;
}
}
if (L[k].size() == 2) {
int mina = mi[L[k][0]];
int minb = mi[L[k][1]];
int maxa = ma[L[k][0]];
int maxb = ma[L[k][1]];
if (maxb < mina) {
res += mina - maxb;
}
if (maxa < minb) {
res += minb - maxa;
}
}
cout << res << endl;
getchar();
getchar();
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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 | #include <iostream> #include <vector> #include <algorithm> using namespace std; const int N = 500001; int H[N], W[N], LV[N], mi[N], ma[N]; vector<int> G[N], L[N]; int n, m; bool vis[N]; int main() { ios_base::sync_with_stdio(0); cin >> n >> m; for (int i = 0; i < n - 1; i++) { int a, b; cin >> a >> b; G[a].push_back(b); G[b].push_back(a); LV[i + 2] = 2 * n; } for (int i = 1; i <= m; i++) { cin >> W[i]; L[0].push_back(i); LV[i] = 0; mi[i] = W[i]; ma[i] = W[i]; } int cnt = m; int k = 0; while (cnt != n) { for (int i = 0; i < L[k].size(); i++) { int leaf = L[k][i]; for (int j = 0; j < G[leaf].size(); j++) { int parent = G[leaf][j]; if (LV[parent] != LV[leaf]) { H[parent]++; if (H[parent] == G[parent].size() - 1) { L[k + 1].push_back(parent); cnt++; LV[parent] = k + 1; } } } } k++; } // cout << "LEVELS" << endl; for (int i = 1; i <= n; i++) { // cout << LV[i] << endl; } long long res = 0; for (int i = 1; i <= k; i++) { for (int j = 0; j < L[i].size(); j++) { int parent = L[i][j]; // cout << "ANALYZING " << parent << endl; vector<int> children_min; vector<int> children_max; vector<long long> children_all; for (int l = 0; l < G[parent].size(); l++) { int leaf = G[parent][l]; if (LV[leaf] >= i) { continue; } children_min.push_back(mi[leaf]); children_max.push_back(ma[leaf]); children_all.push_back(mi[leaf]); children_all.push_back(ma[leaf]); // cout << "ADDING " << mi[leaf] << " " << ma[leaf] << endl; } sort(children_min.begin(), children_min.end()); sort(children_max.begin(), children_max.end()); sort(children_all.begin(), children_all.end()); long long mini = 0, maxi = 0, result = 0, best = 999999999999999999, bestmin, bestmax; long long a = 0, b = 0; for (a = 0; a < children_min.size(); a++) { result += children_min[a] - children_all[0]; } // cout << "STARTING RESULT " << result << endl; a = 0; for (int l = 1; l < children_all.size(); l++) { // cout << "POINT NUMBER " << l << endl; while (a < children_min.size() && children_min[a] < children_all[l]) { a++; // cout << "A UP " << a << endl; } while (b < children_max.size() && children_max[b] < children_all[l]) { b++; // cout << "B UP " << b << endl; } result += b * (children_all[l] - children_all[l - 1]); // cout << "GOING UP BY " << b * (children_all[l] - children_all[l - 1]) << endl; result -= (children_min.size() - a) * (children_all[l] - children_all[l - 1]); // cout << "GOING DOWN BY " << (children_min.size() - a) * (children_all[l] - children_all[l - 1]) << endl; // cout << "CURRENT RESULT " << result << endl; if (result < best) { best = result; bestmin = children_all[l]; bestmax = children_all[l]; } else if (result == best) { bestmax = children_all[l]; } } mi[parent] = bestmin; ma[parent] = bestmax; res += best; // cout << "FINAL RESULT " << res << endl; // cout << mi[parent] << " " << ma[parent] << endl << endl; } } if (L[k].size() == 2) { int mina = mi[L[k][0]]; int minb = mi[L[k][1]]; int maxa = ma[L[k][0]]; int maxb = ma[L[k][1]]; if (maxb < mina) { res += mina - maxb; } if (maxa < minb) { res += minb - maxa; } } cout << res << endl; getchar(); getchar(); return 0; } |
English