#include "roads.h"
#define pb push_back
#define mp make_pair
#include <vector>
#include <set>
#include <algorithm>
using namespace std;
typedef long long ll;
typedef pair<int, ll> pi;
const ll INF = 1e15;
struct SegTree {
int size;
vector<pi> tree;
void init(int N) {
size = 1;
while (size < N) {
size *= 2;
}
tree.assign(2 * size, mp(0, 0));
}
};
int N;
vector<vector<pi>> adj_list;
vector<int> subtree;
vector<int> deg;
// dp1[node][k]: At most k roads open at the subtree of node, road from node to par is either open or closed, whichever is more optimal
vector<vector<ll>> dp1;
// dp2[node][k]: At most k roads open at the subtree of node, road from node to par is closed (IT IS counted in cost)
// dp2 goes up to deg(node), after that it's just the weight
vector<vector<ll>> dp2;
vector<vector<ll>> best1, best2;
void merge(vector<ll>& a, vector<ll>& b) {
if (a.size() < b.size()) {
swap(a, b);
}
for (int i = 0; i < b.size(); i++) {
a[i] = a[i] + b[i];
}
}
bool degcmp(pi a, pi b) {
return deg[a.first] < deg[b.first];
}
inline void insert(multiset<ll>& s, ll& t, ll v, int top) {
t += v;
s.insert(v);
if (top < s.size()) {
t -= *prev(s.end());
s.erase(prev(s.end()));
}
}
ll answer(int target, vector<ll>& diff, multiset<ll>& opt) {
ll ans = INF;
ll cur = 0;
int c1 = 0;
for (int i = 0; i < diff.size(); i++) {
cur += diff[i];
c1++;
}
auto it = opt.begin();
int c2 = 0;
while (c1 + c2 < target && it != opt.end()) {
cur += *it;
c2++;
it++;
}
if (c1 + c2 == target) {
ans = cur;
}
else {
return INF;
}
for (int i = 0; i < diff.size(); i++) {
cur -= diff[i];
c1--;
if (it == opt.end()) {
return ans;
}
cur += *it;
ans = min(ans, cur);
it++;
}
return ans;
}
// Calculate costs coming from children (both for d and d-1 children taken)
void calculate_best_cuts(int node) {
multiset<ll> opt;
ll sum = 0;
best1[node].resize(deg[node] + 1);
best2[node].resize(deg[node] + 1);
sort(adj_list[node].begin(), adj_list[node].end(), degcmp);
int start = 0;
for (int d = 0; d <= deg[node]; d++) {
vector<ll> diff;
int p = start;
while (p < adj_list[node].size()) {
int u;
ll w;
tie(u, w) = adj_list[node][p];
if (deg[u] == d) {
start = p + 1;
// Only need to keep the min deg[node] - d weights
insert(opt, sum, w, deg[node] - d);
}
else {
diff.pb(dp2[u][d] - dp1[u][d]);
}
p++;
}
sort(diff.begin(), diff.end());
diff.resize(min((int)diff.size(), deg[node] - d));
reverse(diff.begin(), diff.end());
while (opt.size() > deg[node] - d) {
opt.erase(prev(opt.end()));
}
// Should have deg[node] - d elements
// If we are to open parent, we should look at closing deg[node] - d kids
// If we are to close parent, we should look at closing deg[node] - d - 1 kids
if (d == 0 && node != 0) {
best1[node][0] = INF;
best2[node][0] = 0;
continue;
}
best1[node][d] = answer(deg[node] - d, diff, opt);
if (d != deg[node]) {
if (diff.size() == deg[node] - d) {
diff.erase(diff.begin());
}
best2[node][d] = answer(deg[node] - d - 1, diff, opt);
}
else {
best2[node][d] = 0;
}
if (node == 0) {
best2[node][d] = INF;
}
}
}
void dfs(int node, int par, ll w) {
deg[node] = adj_list[node].size();
if (par != -1) {
adj_list[node].erase(find(adj_list[node].begin(), adj_list[node].end(), mp(par, w)));
}
subtree[node] = 1;
for (auto [neigh, w] : adj_list[node]) {
dfs(neigh, node, w);
subtree[node] += subtree[neigh];
}
calculate_best_cuts(node);
vector<ll> dp;
for (auto [neigh, w] : adj_list[node]) {
merge(dp1[node], dp1[neigh]);
}
dp2[node].resize(deg[node] + 1);
if (dp1[node].size() < deg[node] + 1) {
dp1[node].resize(deg[node] + 1);
}
for (int d = 0; d <= deg[node]; d++) {
dp2[node][d] = dp1[node][d] + best2[node][d] + w;
dp1[node][d] = min(dp1[node][d] + best1[node][d], dp2[node][d]);
}
}
vector<ll> minimum_closure_costs(int n, vector<int> U, vector<int> V,vector<int> W) {
N = n;
adj_list.resize(N);
dp1.resize(N);
dp2.resize(N);
subtree.resize(N);
deg.resize(N);
best1.resize(N);
best2.resize(N);
for (int i = 0; i < N - 1; i++) {
adj_list[U[i]].pb(mp(V[i], W[i]));
adj_list[V[i]].pb(mp(U[i], W[i]));
}
dfs(0, -1, INF);
while (dp1[0].size() < N) {
dp1[0].pb(0);
}
return dp1[0];
}
