旗鼓相当的对手

NC204272. 旗鼓相当的对手

描述

现在有 $n$ 个炉石玩家被挂在一棵有根树上,树根节点的标号是 $1$ .树是一种有 $n$ 个点且有 $n - 1$ 条边的结构, 并且一棵树保证没有环.

$n$ 个点的每个点上都有一个炉石玩家,并且每个炉石玩家都有一个战力值为 $a_i$ .对于树上的每一位炉石玩家,都有一棵属于他的子树,我们定义两个节点之间的距离为两个节点之间最短路径经过的边数.对于一个给定的值 $k$ ,如果满足 $dis(i, j) = k$ ,我们就说 $i$ 和 $j$ 旗鼓相当. $dis(i,j)$ 表示树上节点 $i$ 通过最短的路径到达节点 $j$ 所经过的边的数量.

现在需要统计每个炉石玩家的 $rating$ 值, $rating$ 值的计算方式是这样的,对于炉石玩家 $u$ 的子树中的所有节点,如果 $x, y$ 是旗鼓相当的,并且 $x,y$ 的最近公共祖先是 $u$ 且满足 $u \neq x \, \& \, u \neq y$ ,那么 $u$ 的 $rating$ 就会增加 $a_x + a_y$ .

输入描述

第一行两个整数,意义和题面给定的一致.
第二行个整数表示战力值.
接下来行每行两个数表示之间有一条边.

输出描述

一行输出个整数表示每个炉石玩家的.按照标号顺序输出.

示例1

输入:

输出:

5 0 0

上次编辑到这里，代码来自缓存点击恢复默认模板

C++14(g++5.4) 解法, 执行用时: 211ms, 内存消耗: 30436K, 提交时间: 2020-04-02 22:16:22

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N = 2e5 + 10;
#define fi first
#define se second
#define pb push_back
int n,f,a[N];
vector<int>v[N];
map<int,LL>s[N];
map<int,LL>sz[N];
LL ans[N];
void dfs(int x,int y){
  LL tot=0;
  for(auto k:v[x]){
    if(k==y)
      continue;
    dfs(k,x);
    for(auto g:s[k]){
      if(g.fi+1<f && sz[x][f-g.fi-1]){
        tot+=g.se*sz[x][f-g.fi-1]+s[x][f-g.fi-1]*sz[k][g.fi];
      }
    }
    for(auto g:s[k]){
      s[x][g.fi+1]+=g.se;
    }
    for(auto g:sz[k]){
      sz[x][g.fi+1]+=g.se;
    }
    sz[k].clear();
    s[k].clear();
  }
  s[x][0]+=a[x];
  sz[x][0]++;
  ans[x]=tot;
}
int main() {
  ios::sync_with_stdio(false);
  cin>>n>>f;
  for(int i=1;i<=n;i++)cin>>a[i];
  for(int i=1;i<n;i++){
    int s,t;
    cin>>s>>t;
    v[s].pb(t);
    v[t].pb(s);
  }
  dfs(1,0);
  for(int i=1;i<=n;i++)cout<<ans[i]<<' ';
  return 0;
}