#include <bits/stdc++.h> using namespace std; struct sol { long long sumRoot, cost; }; sol joinWithEdgeCut(const sol &parent, const sol &child) { return { parent.sumRoot, parent.cost + child.cost }; } sol joinWithEdgeKept(const sol &parent, const sol &child) { return { parent.sumRoot + child.sumRoot, parent.cost + child.cost + 2 * parent.sumRoot * child.sumRoot }; } vector <vector<sol>> dfs(int v, int par, vector <vector<int>> &g, vector <int> &value) { vector <vector<sol>> solParent = {{ {value[v], (long long) value[v] * value[v]} }}; for (int u : g[v]) if (u != par) { auto solChild = dfs(u, v, g, value); vector <vector<sol>> solParentNew(solParent.size() + solChild.size()); for (int cutParent = 0; cutParent < solParent.size(); cutParent++) { for (int cutChild = 0; cutChild < solChild.size(); cutChild++) { int cutTotal = cutParent + cutChild; for (auto &sParent : solParent[cutParent]) { for (auto &sChild : solChild[cutChild]) { solParentNew[cutTotal].push_back(joinWithEdgeKept(sParent, sChild)); solParentNew[cutTotal + 1].push_back(joinWithEdgeCut(sParent, sChild)); } } } } for (auto &sols : solParentNew) { sort(sols.begin(), sols.end(), [&](const sol &a, const sol &b) { if (a.sumRoot != b.sumRoot) { return a.sumRoot < b.sumRoot; } else { return a.cost < b.cost; } }); vector <sol> paretoFront; for (auto &s : sols) { if (paretoFront.empty() || s.cost < paretoFront.back().cost) { paretoFront.push_back(s); } } sols = paretoFront; } solParent = solParentNew; } return solParent; } vector <long long> solve(int n, vector <vector<int>> &g, vector <int> &value) { auto solutions = dfs(0, -1, g, value); vector <long long> bestCosts(n); for (int i = 0; i < n; i++) { bestCosts[i] = solutions[i].back().cost; } return bestCosts; } int main() { ios_base::sync_with_stdio(false); int t; cin >> t; while (t--) { int n; cin >> n; vector <int> value(n); for (int &x : value) { cin >> x; } vector <vector <int>> g(n); for (int e = 1; e < n; e++) { int u, v; cin >> u >> v; u--; v--; g[u].push_back(v); g[v].push_back(u); } auto ans = solve(n, g, value); for (auto &x : ans) { cout << x << ' '; } cout << '\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 111 112 113 114 115 116 117 118 119 120 121 | #include <bits/stdc++.h> using namespace std; struct sol { long long sumRoot, cost; }; sol joinWithEdgeCut(const sol &parent, const sol &child) { return { parent.sumRoot, parent.cost + child.cost }; } sol joinWithEdgeKept(const sol &parent, const sol &child) { return { parent.sumRoot + child.sumRoot, parent.cost + child.cost + 2 * parent.sumRoot * child.sumRoot }; } vector <vector<sol>> dfs(int v, int par, vector <vector<int>> &g, vector <int> &value) { vector <vector<sol>> solParent = {{ {value[v], (long long) value[v] * value[v]} }}; for (int u : g[v]) if (u != par) { auto solChild = dfs(u, v, g, value); vector <vector<sol>> solParentNew(solParent.size() + solChild.size()); for (int cutParent = 0; cutParent < solParent.size(); cutParent++) { for (int cutChild = 0; cutChild < solChild.size(); cutChild++) { int cutTotal = cutParent + cutChild; for (auto &sParent : solParent[cutParent]) { for (auto &sChild : solChild[cutChild]) { solParentNew[cutTotal].push_back(joinWithEdgeKept(sParent, sChild)); solParentNew[cutTotal + 1].push_back(joinWithEdgeCut(sParent, sChild)); } } } } for (auto &sols : solParentNew) { sort(sols.begin(), sols.end(), [&](const sol &a, const sol &b) { if (a.sumRoot != b.sumRoot) { return a.sumRoot < b.sumRoot; } else { return a.cost < b.cost; } }); vector <sol> paretoFront; for (auto &s : sols) { if (paretoFront.empty() || s.cost < paretoFront.back().cost) { paretoFront.push_back(s); } } sols = paretoFront; } solParent = solParentNew; } return solParent; } vector <long long> solve(int n, vector <vector<int>> &g, vector <int> &value) { auto solutions = dfs(0, -1, g, value); vector <long long> bestCosts(n); for (int i = 0; i < n; i++) { bestCosts[i] = solutions[i].back().cost; } return bestCosts; } int main() { ios_base::sync_with_stdio(false); int t; cin >> t; while (t--) { int n; cin >> n; vector <int> value(n); for (int &x : value) { cin >> x; } vector <vector <int>> g(n); for (int e = 1; e < n; e++) { int u, v; cin >> u >> v; u--; v--; g[u].push_back(v); g[v].push_back(u); } auto ans = solve(n, g, value); for (auto &x : ans) { cout << x << ' '; } cout << '\n'; } return 0; } |