牛客练习赛60 旗鼓相当的对手 树上启发式合并

题目链接：https://ac.nowcoder.com/acm/contest/4853/E

#include <bits/stdc++.h>
using namespace std;
#define LL long long
typedef pair<LL, LL> pll;
LL n, k, a[100005];
vector<vector<int> > v(100005);
map<int , pll > mp[100005];//子树到当前点距离 数量 权值
 
LL d[100005];//距离偏移
LL ans[100005];
 
void dfs(int u, int fa){
 
    for(auto to: v[u]){
        if(to==fa) continue;
 
        dfs(to, u); d[to]++;//子树所有点到当前点的距离+1
        if(mp[u].size()<mp[to].size()){//启发式
            swap(mp[u], mp[to]); swap(d[u], d[to]);//swap交换指针O(1)
        }
        //查询
        for(auto y: mp[to]){
            LL p1=y.first+d[to];//子树到当前点的距离
            pll p2=y.second;
 
            if(mp[u].find(k-p1-d[u])!=mp[u].end()){
                pll t=mp[u][k-p1-d[u]];
                ans[u]+=p2.first*t.second+p2.second*t.first;
            }
        }
        //先算答案后合并
        for(auto y: mp[to]){
            LL p1=y.first+d[to];//子树到当前点的距离
 
            pll p2=y.second;
            mp[u][p1-d[u]].first+=p2.first;
            mp[u][p1-d[u]].second+=p2.second;
        }
        mp[to].clear();
    }
    mp[u][-d[u]].first++;//加入u点
    mp[u][-d[u]].second+=a[u];
}
 
int main(){
    scanf("%lld%lld", &n, &k);
    for(int i=1; i<=n; i++){
        scanf("%lld", &a[i]);
    }
    int x, y;
    for(int i=1; i<n; i++){
        scanf("%d%d", &x, &y);
        v[x].push_back(y);
        v[y].push_back(x);
    }
    dfs(1, 0);
    for(int i=1; i<=n; i++){
        printf("%lld ", ans[i]);
    }
 
    return 0;
}