https://ac.nowcoder.com/acm/contest/5278/G
题意很好理解。而且很容易发现树中同一深度的松鼠才会打架。
预处理出节点的深度和dfs序，然后枚举树的深度，同一深度的所有节点和根节点s去建虚树，每建好一次就从根节点s出发跑一次树型dp
dp的转移方程比较好想，设当前节点为x，si为x的子树，则dp[x] = ∑max(1,dp[si] + (depth[si] - depth[x]) ) (dp[si]需要>0)
depth[]表示节点在原树的深度，最终每棵虚树对答案的贡献是 max(1，dp[x] - 1 ）(dp[x]>0)
复杂度O(nlogn)
1 #include<bits/stdc++.h> 2 typedef long long ll; 3 using namespace std; 4 const int maxbit = 20; 5 const int maxn = 2e5+5; 6 vector<int> G[maxn],vt[maxn],p[maxn];//vt为虚树 7 ll a[maxn],dp[maxn]; 8 int depth[maxn],fa[maxn][maxbit],Log[maxn],in[maxn];//in数组为dfs序 9 int n,cnt,s; 10 void add(int u,int v){G[u].push_back(v),G[v].push_back(u);} 11 bool cmp(int u,int v) {return in[u]<in[v];} 12 void pre(){ 13 Log[0] = -1; 14 Log[1] = 0,Log[2] = 1; 15 for(int i = 3;i<maxn;i++) Log[i] = Log[i/2] + 1; 16 } 17 void dfs(int cur,int father){//dfs预处理 18 in[cur] = ++cnt;//处理dfs序 19 depth[cur] = depth[father] + 1;//当前结点的深度为父亲结点+1 20 fa[cur][0] = father;//更新当前结点的父亲结点 21 for(int j = 1;(1<<j)<=n;j++){//倍增更新当前结点的祖先 22 fa[cur][j] = fa[fa[cur][j-1]][j-1]; 23 } 24 for(int i = 0;i<G[cur].size() ;i++){ 25 if(G[cur][i] != father) {//dfs遍历 26 dfs(G[cur][i],cur); 27 } 28 } 29 } 30 int LCA(int u,int v){ 31 if(depth[u]<depth[v]) swap(u,v); 32 int dist = depth[u] - depth[v];//深度差 33 while(depth[u]!=depth[v]){//把较深的结点u倍增到与v高度相等 34 u = fa[u][Log[depth[u]-depth[v]]]; 35 } 36 if(u == v) return u;//如果u倍增到v，说明v是u的LCA 37 for(int i = Log[depth[u]];i>=0;i--){//否则两者同时向上倍增 38 if(fa[u][i]!=fa[v][i]){//如果向上倍增的祖先不同，说明是可以继续倍增 39 u = fa[u][i];//替换两个结点 40 v = fa[v][i]; 41 } 42 } 43 return fa[u][0];//最终结果为u v向上一层就是LCA 44 } 45 46 void build (int indx){//传虚树和深度为indx的数组 47 stack<int> st; 48 st.push(s);//入栈s节点 49 vt[s].clear(); 50 int tmp ; 51 for(int i = 0;i<p[indx].size();i++){ 52 tmp = 0; 53 int cur = LCA(p[indx][i],st.top());//求出栈顶元素和p[i]的LCA 54 while(!st.empty() && LCA(cur,st.top()) != st.top()){//如果LCA和栈顶元素不一样,建栈中虚树的边 55 if(tmp) vt[st.top()].push_back(tmp); 56 tmp = st.top(); 57 st.pop(); 58 } 59 if(st.empty() || st.top()!=cur){ 60 st.push(cur);//如果栈为空或者栈顶元素不等于LCA，入栈 61 vt[cur].clear(); 62 } 63 if(tmp) vt[st.top()].push_back(tmp);//把最后的tmp节点加入链中 64 st.push(p[indx][i]);//当前p[i]入栈 65 vt[p[indx][i]].clear(); 66 } 67 tmp = 0; 68 while(!st.empty()){//加入最后一条链 69 if(tmp) vt[st.top()].push_back(tmp); 70 tmp = st.top(); 71 st.pop(); 72 } 73 } 74 void getdp(int cur){//传需要dp的子虚树 75 dp[cur] = 0; 76 if(vt[cur].size() == 0) { 77 dp[cur] = a[cur]; 78 return; 79 } 80 for(int i = 0;i<vt[cur].size();i++){ 81 int v = vt[cur][i]; 82 getdp(v); 83 if(dp[v]!=0){ 84 dp[cur]+=max((ll)1,dp[v]-(depth[v]-depth[cur])); 85 } 86 } 87 } 88 int main(){ 89 scanf("%d%d",&n,&s); 90 for(int i = 1;i<=n;i++){ 91 scanf("%lld",&a[i]); 92 } 93 for (int i = 0; i < n-1; ++i) 94 { 95 int u,v; 96 scanf("%d%d",&u,&v); 97 add(u,v); 98 /* code */ 99 } 100 pre(); 101 dfs(s,0); 102 ll ans = 0; 103 if(a[s]>1) ans+=(a[s]-1); 104 else ans+=a[s]; 105 for(int i = 1;i<=n;i++) p[depth[i]].push_back(i); 106 for(int i = 2;i<=n;i++){ 107 sort(p[i].begin(),p[i].end(),cmp);//深度一样的节点按dfs序排序 108 if(p[i].size() == 0) continue; 109 build(i);//同一深度的节点建虚树 110 getdp(s); 111 if(dp[s]>1) ans+=dp[s]-1; 112 else ans+=dp[s]; 113 vt[s].clear(); 114 for(int j = 0;j<p[i].size();j++) vt[p[i][j]].clear(); 115 } 116 printf("%lld",ans); 117 return 0; 118 }