Gym100496H-House of Representatives-树

树上每个元素有一个p,元素之间有距离d,计算一个元素u，使得sigma(d(i,u)*pi)最小。

两次dfs，第一次计算本节点以下的sigma()，第二次利用sump求解出ans。

#include <cstdio>
#include <algorithm>
#include <cstring>
#include <vector>
#include <map>

using namespace std;

const int maxn = 100000+10;

int fa[maxn];
long long N,chp[maxn],ans[maxn],chsum[maxn],save_p[maxn],save_up[maxn];
long long sump;
long long final_ans;
int ans_idx;

vector <int> G[maxn];
map<pair<int,int> ,int> save_dis;

void dfs(int u,int p)
{
    int d = G[u].size();
    for(int i=0;i<d;i++)
    {
        int v = G[u][i];
        if(p != v)
        {
            chp[u] += save_p[v];
            dfs(v,fa[v] = u);
            chp[u] += chp[v];
            chsum[u] += (chsum[v]+ save_dis[pair<int,int>(u,v)] *(save_p[v]+chp[v]));
        }
    }
    //printf("index:%d chsum:%d
",u,chsum[u]);
}


void solve(int u,int p)
{
    int d = G[u].size();
    int disp = save_dis[pair<int,int>(u,p)];

    if(p != -1)
    {
        ans[u] =chsum[u] + (ans[p]-(chsum[u]+disp*(save_p[u]+chp[u]))) + disp*(sump-chp[u]-save_p[u]);
        //final_ans = min(ans[u],final_ans);
        //printf("index:%d ans:%d chp:%d dis:%d
",u,ans[u],chp[u],disp);
        if(final_ans > ans[u])
        {
            ans_idx = u;
            final_ans = ans[u];
        }
    }

    for(int i=0;i<d;i++)
    {
        int v = G[u][i];
        if(p != v)
        {
            solve(v,u);
        }
    }
}

int main()
{
    freopen("house.in","r",stdin);
    freopen("house.out","w",stdout);

    scanf("%d",&N);
    for(int i=1;i<=N;i++)
    {
        scanf("%d",&save_p[i]);
        sump += save_p[i];
    }
    for(int i=0;i<N-1;i++)
    {
        int u,v,dis;
        scanf("%d%d%d",&u,&v,&dis);
        G[u].push_back(v);
        G[v].push_back(u);

        save_dis[pair<int,int>(u,v)] = dis;
        save_dis[pair<int,int>(v,u)] = dis;

    }

    int root = 1;
    fa[root] = -1;
    dfs(root,-1);

    ans[root] = chsum[root];
    final_ans = ans[root];
    ans_idx = root;
    solve(root,-1);
    printf("%d %lld
",ans_idx,final_ans);
}