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 = +;

 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=;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\n",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 != -)
     {
         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\n",u,ans[u],chp[u],disp);
         if(final_ans > ans[u])
         {
             ans_idx = u;
             final_ans = ans[u];
         }
     }

     for(int i=;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=;i<=N;i++)
     {
         scanf("%d",&save_p[i]);
         sump += save_p[i];
     }
     for(int i=;i<N-;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 = ;
     fa[root] = -;
     dfs(root,-);

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