P2986 [USACO10MAR]Great Cow Gathering G（换根dp）

Topic portal

This question is an extension of this question .

Question: There is a tree, you can choose a root node, the cost of this root node is $\sum c[i]\ast dis[i]$，$c\left[i\right]$ represents the point weight of point i,$dis\left[i\right]$ represents the distance from point i to the root node.

Idea: First deal with the cost when taking 1 as the root node, and then traverse the tree and transfer according to the equation.

• $f[to]=f[x]-siz\left[i\right]\ast dis(x,to)+(to{t}_{c[i]}-siz[i\left]\right)\ast dis(x,to)$

$siz\left[i\right]$ represents the sum of c[i] in subtree i.

Code:

#include<bits/stdc++.h>
#define endl '\n'
#define mp make_pair
#define pb push_back
#define ll long long
#define int long long
#define pii pair<int,int>
#define sz(x) (int)(x).size()
#define all(x) (x).begin(),(x).end()
#define mem(a,b) memset(a,b,sizeof(a))
char *fs,*ft,buf[1<<20];
#define gc() (fs==ft&&(ft=(fs=buf)+fread(buf,1,1<<20,stdin),fs==ft))?0:*fs++;
inline int read()
{
    
    
    int x=0,f=1;
    char ch=gc();
    while(ch<'0'||ch>'9')
    {
    
    
        if(ch=='-')
            f=-1;
        ch=gc();
    }
    while(ch>='0'&&ch<='9')
    {
    
    
        x=x*10+ch-'0';
        ch=gc();
    }
    return x*f;
}
using namespace std;
const int N=1e6+10;
const int inf=0x3f3f3f3f;
const int mod=1e9+7;
const double eps=1e-7;

vector<pii>e[N];
int siz[N],f[N],n,c[N],m=0;
void dfs1(int fa,int x,int dis)
{
    
    
    siz[x] = c[x];
    for(auto i:e[x])
    {
    
    
        if(i.first==fa)
            continue;
        dfs1(x,i.first,dis+i.second);
        siz[x]+=siz[i.first];
    }
    f[1] = f[1] + c[x]*dis;
}
void dfs2(int fa,int x)
{
    
    
    for(auto i:e[x])
    {
    
    
        if(i.first==fa)
            continue;
        int to = i.first;
        f[to] = f[x]-siz[to]*i.second+(m-siz[to])*i.second;
        dfs2(x,to);
    }
}
void solve()
{
    
    
    cin>>n;
    for(int i=1;i<=n;i++)
        cin>>c[i],m+=c[i];
    for(int i=1;i<n;i++)
    {
    
    
        int u,v,w;
        cin>>u>>v>>w;
        e[u].pb(mp(v,w));
        e[v].pb(mp(u,w));
    }
    dfs1(1,1,0);
    dfs2(1,1);
    cout<<*min_element(f+1,f+n+1)<<endl;
}

signed main()
{
    
    
    solve();

    return 0;
}