Gym100496H-House of Representatives-树

Posted 2020-06-14

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

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

 1 #include <cstdio>
 2 #include <algorithm>
 3 #include <cstring>
 4 #include <vector>
 5 #include <map>
 6 
 7 using namespace std;
 8 
 9 const int maxn = 100000+10;
10 
11 int fa[maxn];
12 long long N,chp[maxn],ans[maxn],chsum[maxn],save_p[maxn],save_up[maxn];
13 long long sump;
14 long long final_ans;
15 int ans_idx;
16 
17 vector <int> G[maxn];
18 map<pair<int,int> ,int> save_dis;
19 
20 void dfs(int u,int p)
21 {
22     int d = G[u].size();
23     for(int i=0;i<d;i++)
24     {
25         int v = G[u][i];
26         if(p != v)
27         {
28             chp[u] += save_p[v];
29             dfs(v,fa[v] = u);
30             chp[u] += chp[v];
31             chsum[u] += (chsum[v]+ save_dis[pair<int,int>(u,v)] *(save_p[v]+chp[v]));
32         }
33     }
34     //printf("index:%d chsum:%d\n",u,chsum[u]);
35 }
36 
37 
38 void solve(int u,int p)
39 {
40     int d = G[u].size();
41     int disp = save_dis[pair<int,int>(u,p)];
42 
43     if(p != -1)
44     {
45         ans[u] =chsum[u] + (ans[p]-(chsum[u]+disp*(save_p[u]+chp[u]))) + disp*(sump-chp[u]-save_p[u]);
46         //final_ans = min(ans[u],final_ans);
47         //printf("index:%d ans:%d chp:%d dis:%d\n",u,ans[u],chp[u],disp);
48         if(final_ans > ans[u])
49         {
50             ans_idx = u;
51             final_ans = ans[u];
52         }
53     }
54 
55     for(int i=0;i<d;i++)
56     {
57         int v = G[u][i];
58         if(p != v)
59         {
60             solve(v,u);
61         }
62     }
63 }
64 
65 int main()
66 {
67     freopen("house.in","r",stdin);
68     freopen("house.out","w",stdout);
69 
70     scanf("%d",&N);
71     for(int i=1;i<=N;i++)
72     {
73         scanf("%d",&save_p[i]);
74         sump += save_p[i];
75     }
76     for(int i=0;i<N-1;i++)
77     {
78         int u,v,dis;
79         scanf("%d%d%d",&u,&v,&dis);
80         G[u].push_back(v);
81         G[v].push_back(u);
82 
83         save_dis[pair<int,int>(u,v)] = dis;
84         save_dis[pair<int,int>(v,u)] = dis;
85 
86     }
87 
88     int root = 1;
89     fa[root] = -1;
90     dfs(root,-1);
91 
92     ans[root] = chsum[root];
93     final_ans = ans[root];
94     ans_idx = root;
95     solve(root,-1);
96     printf("%d %lld\n",ans_idx,final_ans);
97 }