思路为树的重心,一遍dfs找重心,二遍dfs找答案,但后面神奇报零,求大佬解答。(样例能过)
#include <iostream>
#include <vector>
#include <cstring>
#define int long long
using namespace std;
const int N = 1e5 + 5;
struct qq{
int u,v,w;
}d[N];
int n,c[N],f[N],siz[N],ans,maxsiz[N],s,dis[N],sum;
bool f1;
vector <pair<int,int> > g[N];
void add (int u,int v,int w){
g[u].push_back({v,w});
g[v].push_back({u,w});
}
void dfs (int x,int father){
// f[x] = c[x];
siz[x] = c[x];
f1 = 0;
for (auto y:g[x]){
int v = y.first;
// int w = y.second;
if (v == father) continue;
dfs (v,x);
siz[x] += siz[v];
// ans[x] = siz[v] * f[v];
if (siz[v] * 2 > sum) f1 = 1;
}
if (siz[x] * 2 < sum) f1 = 1;
// maxsiz[x] = max (maxsiz[x],siz[v]);
if (!f1){
s = x;
}
}
void dfs1 (int x,int father){
siz[x] = c[x];
for (auto y:g[x]){
int v = y.first;
int w = y.second;
if (v == father) continue;
dfs1 (v,x);
siz[x] += siz[v];
ans += 1ll * siz[v] * w;
}
}
signed main(){
cin >> n;
for (int i = 1;i <= n;i ++){
cin >> c[i];
sum += c[i];
}
for (int i = 1;i <= n - 1;i ++){
cin >> d[i].u >> d[i].v >> d[i].w;
add (d[i].u,d[i].v,d[i].w);
}
dfs (1,0);
// for (int i = 1;i <= n;i ++){
//
// if (maxsiz[i] < maxsiz[id]){
//
// id = i;
//
// }
//
// }
// dis[id] = 0;
memset(siz,0,sizeof(siz));
dfs1(s,0);
// for (int i = 1;i <= n;i ++){
//
// cout << ans[i] << '\n';
//
// }
cout << ans << '\n';
return 0;
}