#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; } |