套用长链剖分板子...
复杂度为
#include<bits/stdc++.h> using namespace std; typedef double db; typedef long long ll; typedef unsigned long long ull; typedef pair<int, int> pii; #define fi first #define se second #define empb emplace_back #define all(x) (x).begin(),(x).end() const int N = 1e5 + 100; struct node_t { ll a[2]; const ll &operator [](int i)const { return a[i];} ll &operator [] (int i) { return a[i];} void init(ll x) { a[0] = 1; a[1] = x; } node_t& operator += (const node_t &rhs) { a[0] += rhs[0], a[1] += rhs[1]; return *this; } }; vector<int> E[N]; int n, k; int a[N], dep[N], son[N]; void prepare(int u, int f) { for(auto &v : E[u]) { if(v == f) continue; prepare(v, u); if(dep[son[u]] < dep[v]) son[u] = v; } dep[u] = dep[son[u]] + 1; } node_t pool[N], *pt = pool, *dp[N]; // dp[i][j]表示以i为根的子树下距离i为j的点数以及权值和... // 长链直接继承长儿子,短链暴力转移... // 对于点u处的答案为\sum_{v} dp[u][i][0]*dp[v][k-i-1][1]+dp[u][i][1]*dp[v][i][0] ll res[N]; void dfs(int u, int f) { dp[u] = pt++; dp[u]->init(a[u]); if(son[u]) dfs(son[u], u); for(auto &v : E[u]) { if(v == son[u] || v == f) continue; dfs(v, u); for(int l = 0; l <= dep[v] && l < k - 1; l++) { if(k - l - 1 > dep[u]) continue; res[u] += dp[v][l][0] * dp[u][k - l - 1][1]; res[u] += dp[v][l][1] * dp[u][k - l - 1][0]; } for(int l = 0; l <= dep[v]; l++) { dp[u][l + 1] += dp[v][l]; } } } int main() { #ifdef local freopen("in.txt", "r", stdin); #endif ios::sync_with_stdio(false); cin.tie(0), cout.tie(0); cin >> n >> k; for(int i = 1; i <= n; i++) cin >> a[i]; for(int i = 1, u, v; i < n; i++) { cin >> u >> v; E[u].push_back(v); E[v].push_back(u); } dep[0] = -1; prepare(1, 0); dfs(1, 0); for(int i = 1; i <= n; i++) { cout << res[i] << " \n"[i == n]; } return 0; }