HJ135 计树

知识点动态规划

校招时部分企业笔试将禁止编程题跳出页面，为提前适应，练习时请使用在线自测，而非本地IDE。

描述

给定由 nn 个结点构成、以 11 号节点为根节点的有根树，选中其中 kk 个节点，记为集合 VV。

现在，你需要构建一个计数数组 cntcnt，其中 cnticnti​ 表示节点 ii 作为 LCA 的次数。

具体操作如下：

∙ ∙从集合 VV 中选择一个节点 uu；

∙ ∙从集合 VV 中选择一个节点 vv（可能会与 uu 相同）；

∙ ∙记 u,vu,v 两个节点的最近公共祖先（LCA）为 ii，更新 cnticnti​ 为 cnti+1cnti​+1。

对于全部 k2k2 个选取方式，重复上述操作。最后输出 cntcnt 数组。

输入描述：

第一行输入一个整数 n(1≦n≦105)n(1≦n≦105) 代表树的节点数。

此后 n−1n−1 行，第 ii 行输入两个整数 ui,vi(1≦ui,vi≦n; ui≠vi)ui​,vi​(1≦ui​,vi​≦n; ui​​=vi​) 代表树上第 ii 条边连接节点 uiui​ 和 vivi​。

第 n+1n+1 行输入一个整数 k(1≦k≦n)k(1≦k≦n) 代表集合 VV 的大小。

第 n+2n+2 行输入 kk 个整数 a1,a2,...,ak(1≦ai≦n)a1​,a2​,...,ak​(1≦ai​≦n) 代表集合 VV 中的节点。

输出描述：

在一行上输出 nn 个整数，其中第 ii 个整数表示节点 ii 作为 LCA 的次数，即 cnticnti​ 的值。

示例1

输入：

5
1 2
1 3
2 4
2 5
3
2 3 4

复制输出：

4 3 1 1 0

复制说明：

在这个样例中，树的形态如下图所示。

#include <iostream>
#include <vector>

using namespace std;

const int MAXN = 100005;
vector<vector<int>> adj; // 邻接表
bool isInV[MAXN] = {false}; // 统计节点是否属于集合V
long long countV[MAXN] = {0};  // countV[u]统计以u为根的子树中含有V节点的数量
long long LcaCount[MAXN] = {0}; // Final answer cnt[u]

// DFS function to calculate countV and LcaCount
long long dfs(int u, int p) { // 
    long long current_v_num = isInV[u] == true ? 1 : 0;
    long long child_v_size = 0;
    for(auto v: adj[u]){
        if(v != p){ // 只往孩子节点方向遍历，防止往回遍历到父节点
            long long child_v_num = dfs(v, u);
            current_v_num += child_v_num;
            child_v_size += child_v_num*child_v_num;
        }
    }
    LcaCount[u] = current_v_num * current_v_num - child_v_size; // 记录LCA值
    countV[u] = current_v_num;
    return current_v_num;
}

int main() {
    ios_base::sync_with_stdio(false); // Faster I/O
    cin.tie(NULL);
    
    int n, u, v, k;
    cin >> n;
    adj.resize(n + 1);
    for(int i=0; i<n-1; i++){
        cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }
    cin >> k;
    for(int i=0; i<k; i++){
        cin >> u;
        isInV[u] = true;
    }
    dfs(1, 0); // 从根节点开始遍历，根节点的父节点设为0
    for(int i=1; i<=n; i++){
        cout << LcaCount[i] << (i==n ? "" : " ");
    }
}