题目链接: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; }