Codeforces 1092F Tree with Maximum Cost(树形DP)
题意:给定一棵树,树上每个顶点都有属性值ai,树的边权为1,求$\sum\limits_{i = 1}^{n} dist(i, v) \cdot a_i$,$dist(i, v) $为顶点i到顶点v的距离。该顶点v可以任意选择。
题解:O(n^2)的做法:从每个顶点跑一遍DFS,计算贡献值,并更新答案。(超时)
我们可以先计算出从顶点1跑的答案,发现顶点之间贡献的转移为$ans[u]=ans[fa]+(all-sum[u])-sum[u]$。(all为$\sum\limits_{i = 1}^{n} a_i$)
该顶点的上半部分贡献值增加(all-sum[u]),下半部分贡献值减少(sum[u])。
1 #include <set> 2 #include <map> 3 #include <queue> 4 #include <deque> 5 #include <stack> 6 #include <cmath> 7 #include <cstdio> 8 #include <vector> 9 #include <string> 10 #include <cstring> 11 #include <fstream> 12 #include <iostream> 13 #include <algorithm> 14 using namespace std; 15 16 #define eps 1e-8 17 #define pb push_back 18 #define PI acos(-1.0) 19 #define INF 0x3f3f3f3f 20 #define clr(a,b) memset(a,b,sizeof(a) 21 #define bugc(_) cerr << (#_) << " = " << (_) << endl 22 #define FAST_IO ios::sync_with_stdio(false);cin.tie(NULL);cout.tie(NULL) 23 24 const int N=2e5+10; 25 typedef long long ll; 26 typedef unsigned long long ull; 27 ll a[N],sum[N],res,cnt,ans[N]; 28 29 vector <int> E[N]; 30 31 void dfs(int u,int fa,ll len){ 32 res+=len*a[u]; 33 sum[u]=a[u]; 34 for(int i=0;i<E[u].size();i++){ 35 int v=E[u][i]; 36 if(v==fa) continue; 37 dfs(v,u,len+1); 38 sum[u]+=sum[v]; 39 } 40 } 41 42 void DFS(int u,int fa){ 43 if(fa!=0) ans[u]=ans[fa]+cnt-2*sum[u]; 44 for(int i=0;i<E[u].size();i++){ 45 int v=E[u][i]; 46 if(v==fa) continue; 47 DFS(v,u); 48 } 49 } 50 51 int main(){ 52 FAST_IO; 53 int n; 54 cin>>n; 55 for(int i=1;i<=n;i++) cin>>a[i],cnt+=a[i]; 56 for(int i=1;i<n;i++){ 57 int u,v; 58 cin>>u>>v; 59 E[u].push_back(v); 60 E[v].push_back(u); 61 } 62 dfs(1,0,0); 63 ans[1]=res; 64 DFS(1,0); 65 cout<<*max_element(ans+1,ans+1+n)<<endl; 66 return 0; 67 }