题解 | 计树

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

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int n;
    cin >> n;
    vector<vector<int>> edges(n + 1);
    for (int i = 0; i < n - 1; ++i) {
        int u, v;
        cin >> u >> v;
        edges[u].push_back(v);
        edges[v].push_back(u);
    }

    // 构建子节点列表
    vector<int> parent(n + 1, 0);
    vector<vector<int>> children(n + 1);
    queue<int> q;
    q.push(1);
    parent[1] = 0;
    while (!q.empty()) {
        int u = q.front();
        q.pop();
        for (int v : edges[u]) {
            if (v != parent[u]) {
                parent[v] = u;
                children[u].push_back(v);
                q.push(v);
            }
        }
    }

    int k;
    cin >> k;
    vector<int> V(k);
    for (int i = 0; i < k; ++i) {
        cin >> V[i];
    }

    // 标记是否在集合 V 中
    vector<int> is_in_V(n + 1, 0);
    for (int node : V) {
        is_in_V[node] = 1;
    }

    // 计算 s 数组
    vector<long long> s(n + 1, 0);
    vector<pair<int, bool>> stack;
    stack.emplace_back(1, false);
    while (!stack.empty()) {
        auto [node, visited] = stack.back();
        stack.pop_back();
        if (!visited) {
            stack.emplace_back(node, true);
            // 逆序压入子节点以保证处理顺序正确
            for (auto it = children[node].rbegin(); it != children[node].rend(); ++it) {
                stack.emplace_back(*it, false);
            }
        } else {
            long long s_val = is_in_V[node];
            for (int child : children[node]) {
                s_val += s[child];
            }
            s[node] = s_val;
        }
    }

    // 计算 cnt 数组
    vector<long long> cnt(n + 1, 0);
    for (int i = 1; i <= n; ++i) {
        long long total = s[i] * s[i];
        long long sum_child = 0;
        for (int child : children[i]) {
            sum_child += s[child] * s[child];
        }
        cnt[i] = total - sum_child;
    }

    // 输出结果
    for (int i = 1; i <= n; ++i) {
        cout << cnt[i] << ' ';
    }

    return 0;
}