【专题训练】医院设置[2] | 树的边、点都带权重心

【问题描述】
阆中市管辖了n个乡镇（编号为1~n），有n-1条公路把他们连接起来，每条公路的长度均为c[i]。其中第i个乡镇有人口数量为Wi。

市政府打算在这n个乡镇中选择一个建立一所医院，使所有镇民到医院的所走路程总和最小。

例如，下图中有5个乡镇，每个圈表示一个乡镇，圈外的数字表示该乡镇编号，圈里的数字表示该乡镇的人口数量。
医院应建立在乡镇3，所有人走的路程总和为：13*1+4*2+20*1+40*1=81

【输入格式】
第一行包含 1 个整数 n，表示乡镇数量（编号为1~n）。
接下来 n-1 行，每行包含 2 个用空格隔开的正整数 u,v,c，表示编号为 u 和编号为 v 的乡镇之间有一条公路相连，公路长度为c。
最后 1 行，包含 n 个正整数，每两个正整数之间用一个空格隔开，其中第 i 个整数表示编号为 i 的乡镇的人口数量为 Wi

【输出格式】
第一行包含若干整数，表示可以建立医院的乡镇编号（由小到达输出）。第二行一个整数，表示所有人行程总和。

【输入样例】

5
1 2 1
1 3 1
3 4 1
3 5 1
13 4 12 20 40

【输出样例】

3
81

【数据范围】
对于 30%的数据， 1 < n ≤ 100；
对于 60%的数据， 1 < n ≤ 2000；
对于 100%的数据， 1 < n ≤ 200 ,000，0 < ci,wi ≤ 10 ,000。

代码：

#include<cstdio>
#include<vector>
#include<iostream>
#include<cstring>
#define maxn 200005
using namespace std;

int n,popu[maxn],maxsum[maxn],tot=0;
vector<int>g[maxn],w[maxn];

int read()
{
    int x=0;
    char ch=getchar();
    while(ch<'0' || ch>'9') ch=getchar();
    while(ch>='0' && ch<='9')
    {
        x=x*10+ch-'0';
        ch=getchar();
    }
    return x;
}

void init()
{
    n=read();
    int x,y,z;
    for(int i=1;i<n;i++)
    {
        x=read();
        y=read();
        z=read();
        g[x].push_back(y);
        g[y].push_back(x);
        w[x].push_back(z);
        w[y].push_back(z);
    }
    for(int i=1;i<=n;i++)
    {
        popu[i]=read();
        tot+=popu[i];
    }
}

bool vis[maxn];
int sum[maxn];
long long cost[maxn];

void DFS(int i)
{
    vis[i]=true;
    sum[i]=popu[i];
    cost[i]=0;
    maxsum[i]=0;
    for(int k=0;k<g[i].size();k++)
    {
        int j=g[i][k],c=w[i][k];
        if(vis[j]) continue;
        DFS(j);
        sum[i]+=sum[j];
        maxsum[i]=max(maxsum[i],sum[j]);
        cost[i]+=cost[j]+sum[j]*c;
    }
    maxsum[i]=max(maxsum[i],tot-sum[i]);
} 

void solve()
{
    init();
    DFS(1);
    int sum_min=tot,x;
    for(int i=1;i<=n;i++)
    {
        if(maxsum[i]<sum_min)
        {
            sum_min=maxsum[i];
            x=i;
        }
    }
    for(int i=1;i<=n;i++)
    {
        if(maxsum[i]==sum_min)
        {
            printf("%d ",i);
        }
    }
    printf("\n");
    memset(vis,0,sizeof(vis));
    DFS(x);
    cout<<cost[x]<<endl;
}

int main()
{
    //freopen("in.txt","r",stdin);
    solve();
    return 0;
}