牛客 - 练习赛60E 旗鼓相当的对手

传送门：牛客 - 练习赛60E

思路分析

子树问题启发式合并，保留重儿子的贡献，删除轻儿子的贡献，暴力合并轻儿子的贡献
开两个map，num表示该深度的数量，sum表示该深度的点权和

样例输入

3 2
1 2 3
1 2
1 3

样例输出

5 0 0

AC代码

#include <functional>
#include <algorithm>
#include <iostream>
#include <cstring>
#include <cstdlib>
#include <iomanip>
#include <vector>
#include <string>
#include <cstdio>
#include <chrono>
#include <random>
#include <queue>
#include <stack>
#include <cmath>
#include <map>
#include <set>
#if __cplusplus >= 201103L
#include <unordered_map>
#include <unordered_set>
#endif
#define ls x<<1
#define rs x<<1|1
#define fi first
#define se second
#define ll long long
#define pb push_back
#define mp make_pair
#define fun function
#define sz(x) x.size()
#define lowbit(x) x&(-x)
#define all(x) x.begin(),x.end()
#define mem(a,b) memset(a,b,sizeof(a))
#define IOS ios_base::sync_with_stdio(0); cin.tie(0);cout.tie(0);

using namespace std;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

const int INF = 0x3f3f3f3f;
const int N=1e5+10;

int n,k,siz[N],son[N],dep[N];
vector<int>e[N];
vector<pair<int,ll>>tem;
ll ans[N],a[N];
map<int,ll>num,sum;

void dfs1(int u,int fa) {
    siz[u]=1;
    dep[u]=dep[fa]+1;
    for(auto v:e[u]) {
        if(v==fa) continue;
        dfs1(v,u);
        siz[u]+=siz[v];
        if(siz[v]>siz[son[u]]) son[u]=v;
    }
}

void add(int u,int fa,int pos) {
    ans[pos]+=num[k+2*dep[pos]-dep[u]]*a[u]+sum[k+2*dep[pos]-dep[u]];
    tem.push_back(mp(dep[u],a[u]));
    for(auto v:e[u]) if(v!=fa) add(v,u,pos);
}

void del(int u,int fa) {
    num[dep[u]]--;
    sum[dep[u]]-=a[u];
    for(auto v:e[u]) if(v!=fa) del(v,u);
}

void dfs2(int u,int fa) {
    for(auto v:e[u]) {
        if(v==fa||v==son[u]) continue;
        dfs2(v,u);
        del(v,u);
    }
    if(son[u]) dfs2(son[u],u);
    for(auto v:e[u]) {
        if(v==fa||v==son[u]) continue;
        tem.clear();
        add(v,u,u);
        for(auto x:tem) {
            num[x.fi]++;
            sum[x.fi]+=x.se;
        }
    }
    num[dep[u]]++;
    sum[dep[u]]+=a[u];
}


int main() {
    IOS;
#ifdef xiaofan
    freopen("1.in","r",stdin);
    freopen("1.out","w",stdout);
#endif

    cin>>n>>k;
    for(int i=1; i<=n; i++) cin>>a[i];
    for(int i=1; i<n; i++) {
        int u,v;
        cin>>u>>v;
        e[u].push_back(v);
        e[v].push_back(u);
    }
    dfs1(1,1);
    dfs2(1,1);
    for(int i=1; i<=n; i++) cout<<ans[i]<<" ";



    return 0;
}