题解 | 在树上游玩

#include <iostream>
#include <vector>
using namespace std;
typedef long long LL;
const int MOD = 1e9 + 7;
const int maxn = 200005;
bool vis[maxn] ={false}; // 标记数组，记录节点是否被访问过
bool masked[maxn] = {false}; // 标记数组，记录节点是否为被标记节点
vector<int> G[maxn]; // 邻接表，存储树的结构

// 统计v能到达的neighbor数量
void DFS(int v, int &cnt){
    if(vis[v]) return; // 如果节点v已经被访问过，直接返回
    vis[v] = true; // 标记节点v为已访问
    for(int i=0; i<G[v].size(); i++){ // 遍历节点v的所有邻接点
        if(masked[G[v][i]] && vis[G[v][i]] == false){ // 若邻接点是标记节点且未被访问
            DFS(G[v][i], cnt); // 继续深度优先搜索
        }else if(vis[G[v][i]] == false){ // 若邻接点未被访问且不是标记节点
            cnt++; // neighbor数量加1
        }
    }
}

int main() {
    int n, k, u, v;
    cin >> n >> k; // 输入树的节点数和被标记的点数
    for(int i=0; i<k; i++){ // 输入被标记的点
        cin >> v;
        masked[v] = true; // 标记节点
    }
    for(int i=0; i<n-1; i++){ // 输入树的边
        cin >> u >> v;
        G[u].push_back(v); // 添加边
        G[v].push_back(u);
    }
    LL ans = 1, cost = 0; // ans存储染色方案数量，cost存储最小染色代价
    for(int i=1; i<=n; i++){ // 遍历所有节点
        if(masked[i] && vis[i] == false){ // 若节点是标记节点且未被访问
            int cnt = 0; // 初始化neighbor数量
            DFS(i, cnt); // 统计neighbor数量
            cost++; // 最小代价为标记节点连通块数量
            ans = (ans * cnt) % MOD; // 方案数为当前方案数乘上每个连通块的邻居节点数量
        }
    }
    cout << cost << ' ' << ans; // 输出结果
    return 0;
}