小松鼠的聚会(Squirrel Gathering) 题解

思路分析

按题目要求使用对应数据结构与算法模拟/递推求解，注意数据范围边界与时间复杂度。

参考代码

#include <iostream>
#include <vector>
using namespace std;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    int n;
    if (!(cin >> n)) return 0;
    vector<long long> c(n + 1);
    for (int i = 1; i <= n; ++i) cin >> c[i];
    vector<vector<int>> adj(n + 1);
    for (int i = 0; i < n - 1; ++i) {
        int u, v;
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    int k;
    cin >> k;
    vector<int> par(n + 1, 0), order;
    order.reserve(n);
    order.push_back(1);
    for (size_t head = 0; head < order.size(); ++head) {
        int u = order[head];
        for (int v : adj[u]) {
            if (v == par[u]) continue;
            par[v] = u;
            order.push_back(v);
        }
    }
    vector<vector<long long>> down(n + 1, vector<long long>(k + 1, 0));
    for (int i = n - 1; i >= 0; --i) {
        int u = order[i];
        down[u][0] = c[u];
        for (int v : adj[u]) {
            if (v == par[u]) continue;
            for (int d = 1; d <= k; ++d) {
                down[u][d] += down[v][d - 1];
            }
        }
    }
    vector<vector<long long>> tot = down;
    for (int u : order) {
        for (int v : adj[u]) {
            if (v == par[u]) continue;
            for (int d = k; d >= 1; --d) {
                long long from_u = tot[u][d - 1] - (d >= 2 ? down[v][d - 2] : 0LL);
                tot[v][d] += from_u;
            }
        }
    }
    for (int i = 1; i <= n; ++i) {
        long long sum = 0;
        for (int d = 0; d <= k; ++d) sum += tot[i][d];
        cout << sum << "\n";
    }
    return 0;
}