牛客竞赛题目讲解_旗鼓相当的对手(树上启发式合并)

// https://ac.nowcoder.com/acm/contest/4853/E

#include "stdafx.h"

//#include <bits/stdc++.h>

#include <iostream>

#include <cstdio>

#include <cstring>

#include <algorithm>

 

#include <vector> 

 

using namespace std; 

const int maxn = 1e5 + 5;

 

int n,k, ac[maxn]={0,

1,2,3,4,5,6,7,9 

};

int edge[32][2]={

1 ,2,

1, 3,

3,4,

1,5,

3,6,

5,7,

5,8

};

  

long long ans[maxn], sum[maxn]; 

int sz[maxn], son[maxn],dep[maxn],cnt[maxn];

int dfsn[maxn],T=0;

int a[maxn];//dfs序 编号对应的原标号

vector<int>load[maxn];

 

void dfs1(int s, int pre) {

    sz[s] = 1;// 子树大小

dep[s] = dep[pre] + 1;// 节点深度

  

dfsn[s]=++T; //dfs序 编号

a[T]=s; //dfs序 编号对应的原标号

    for(auto e : load[s]) {

        if(e == pre)   continue;

        dfs1(e, s);

        sz[s] += sz[e];

        if(sz[e] > sz[son[s]])

            son[s] = e;// 重子树

    }

}

void dfs2(int s, int pre, int flag) {

    for(auto e : load[s]) {

        if(e != pre && e != son[s])// 先 计算轻链

        dfs2(e, s, 0);

    }

    if(son[s])        dfs2(son[s], s, 1);// 最后计算重链，而不需要清除 cnt[maxn];

    

for(auto e : load[s]) 

   if(e != pre && e != son[s])// 只 计算轻链

{  

  for(int i = 0; i <sz[e]; i++) 

  {

int dis = k + 2 * dep[s] - dep[a[dfsn[e]+i]];

// cout<<" pre="<<pre<<", ans[pre] ="<<ans[pre] <<", s="<<s<<", a[s] ="<<a[s] <<", dep[s]="<<dep[s]<<", cnt[dis]="<<cnt[dis]<<", dis="<<dis<<", sum[dis] ="<<sum[dis] <<endl;

if(dis  >dep[s]) ans[s] += sum[dis] + ac[a[dfsn[e]+i]] * cnt[dis]; 

cout<<" s="<<s<<", ans[s] ="<<ans[s] <<", a[dfsn[e]+i]="<<a[dfsn[e]+i]<<", ac[a[dfsn[e]+i]] ="<<ac[a[dfsn[e]+i]] <<", dep[a[dfsn[e]+i]]="<<dep[a[dfsn[e]+i]]<<", cnt[dis]="<<cnt[dis]<<", dis="<<dis<<", sum[dis] ="<<sum[dis] <<endl;

  }

    for(int i = 0; i <sz[e]; i++) //  对一个分支统计完贡献后，再添加它的信息cnt[dep]和sum[dep]

  {

sum[dep[a[dfsn[e]+i]]] +=  ac[a[dfsn[e]+i]];

cnt[dep[a[dfsn[e]+i]]] += 1;    

  }

}

sum[dep[s]] += ac[s] , cnt[dep[s]] ++ ;// 统计根节点信息cnt[dep]和sum[dep]

   if(!flag) 

{// 轻链 clear

         for(int i = 0; i <sz[s]; i++) {

sum[dep[a[dfsn[s]+i]]] -=  ac[a[dfsn[s]+i]];

cnt[dep[a[dfsn[s]+i]]] -= 1;    

  } 

    }

}

int main() 

{

n=8;k=3;

    //cin >> n>>k  

     

        for(int i = 1; i <= n; i++) {

          //  cin >> ac[i];

            

        }

int s,e;

        for(int i = 1; i < n; i++) {

            //cin >> s >> e;

s=edge[i-1][0];

e=edge[i-1][1];

            load[s].push_back(e);

            load[e].push_back(s);

        }

        dfs1(1, 0);

  for(int i = 1; i <=n; i++) cout<<"i=  "<<i<<":   son[i]="<<son[i]<<endl;

        dfs2(1, 0, 1);

        for(int i = 1; i <=n; i++) {          

                cout << ans[i]<< " \n"[i == n];          

        }    

}