题解 | 计树

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

const int MAXN = 1e5 + 5;

// 函数：构建子节点列表和父节点数组
void buildChildrenAndParent(int n, const vector<vector<int>>& edges, vector<int>& parent, vector<vector<int>>& children) {
    queue<int> q;
    q.push(1); // 从根节点1开始
    parent[1] = 0; // 根节点的父节点是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); // 将邻接点加入队列，以便后续处理
            }
        }
    }
}

// 函数：计算 s 数组
void calculateSArray(int n, const vector<vector<int>>& children, const vector<int>& is_in_V, vector<long long>& s) {
    vector<pair<int, bool>> stack; // 用于深度优先搜索（DFS）的栈
    stack.emplace_back(1, false); // 从根节点1开始，标记为未访问
    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]; // 当前节点的s值初始化为是否在集合V中
            for (int child : children[node]) { // 遍历当前节点的子节点
                s_val += s[child]; // 累加子节点的s值
            }
            s[node] = s_val; // 更新当前节点的s值
        }
    }
}

// 函数：计算 cnt 数组
void calculateCntArray(int n, const vector<vector<int>>& children, const vector<long long>& s, vector<long long>& cnt) {
    for (int i = 1; i <= n; ++i) { // 遍历所有节点
        long long total = s[i] * s[i]; // 当前节点的s值的平方，表示所有可能的节点对
        long long sum_child = 0; // 子节点的s值的平方和
        for (int child : children[i]) { // 遍历当前节点的子节点
            sum_child += s[child] * s[child]; // 累加子节点的s值的平方
        }
        cnt[i] = total - sum_child; // 当前节点作为LCA的次数等于总对数减去子节点的对数
    }
}

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);
    }

    int k; // 集合V的大小
    cin >> k;
    vector<int> V(k); // 集合V中的节点
    for (int i = 0; i < k; ++i) {
        cin >> V[i]; // 读取集合V中的节点
    }

    // 标记是否在集合 V 中
    vector<int> is_in_V(n + 1, 0); // 标记数组，记录每个节点是否在集合V中
    for (int node : V) {
        is_in_V[node] = 1; // 如果节点在集合V中，标记为1
    }

    vector<int> parent(n + 1, 0); // 存储每个节点的父节点
    vector<vector<int>> children(n + 1); // 存储每个节点的子节点
    buildChildrenAndParent(n, edges, parent, children); // 构建子节点列表和父节点数组

    vector<long long> s(n + 1, 0); // s[i] 表示以节点i为根的子树中，集合V中的节点数
    calculateSArray(n, children, is_in_V, s); // 计算 s 数组

    vector<long long> cnt(n + 1, 0); // cnt[i] 表示节点i作为LCA的次数
    calculateCntArray(n, children, s, cnt); // 计算 cnt 数组

    // 输出结果
    for (int i = 1; i <= n; ++i) {
        cout << cnt[i] << ' '; // 输出每个节点作为LCA的次数
    }
    cout << endl;

    return 0;
}