题解 | 在树上游玩

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

int main() {
    // 
    int n,k;
    cin >> n >> k;

    // denote the red nodes among the n nodes
    vector<bool> red_nodes_tag(n+1,false);
    vector<int> red_nodes;
    for(int i=0;i<k;i++){
        int v;
        cin >> v;
        red_nodes.push_back(v);
        red_nodes_tag[v]=true;
    }

    // get the tree by adjacent matrix
    vector<vector<int>> tree(n+1);
    int u,v;
    while(cin >> u >> v){
        tree[u].push_back(v);
        tree[v].push_back(u);
    }


    // use dfs to populate to get connected part, by the way cnt the num of not red neighbor
    vector<int> already_dfs(n+1,0);
    vector<int> visited_as_root(n+1,0);
    
    auto dfs = [&already_dfs,&visited_as_root,&red_nodes_tag,&tree](auto self, int root)->int{
        already_dfs[root] = 1;

        if(visited_as_root[root] && red_nodes_tag[root]) return 0;
        if(!red_nodes_tag[root]) return 1; // impossible case

        long long neighbor_cnt = 0;
        for(auto v: tree[root]){
            if(!already_dfs[v] && red_nodes_tag[v]){
                neighbor_cnt += self(self, v);
                // cout << "v" << v << ';' << neighbor_cnt << endl;
            }
            else if(!red_nodes_tag[v]){
                // visited_as_leaf[v] = 1;
                neighbor_cnt += 1;
            }
        }
        visited_as_root[root] = 1;
        return neighbor_cnt;
    };


    int part_cnt = 0;
    long long num_methods = 1;
    long long base = 1e9+7;
    
    for(auto v: red_nodes){
        if(!visited_as_root[v]){
            part_cnt += 1;
            long long neighbor_cnt=0;
            neighbor_cnt = dfs(dfs, v);
            // cout << "dfs" << ' ' << v << "neighbor_cnt" << neighbor_cnt << endl;
            num_methods = ((num_methods % base) * (neighbor_cnt % base)) % base;
            // cout << neighbor_cnt << ' ' << num_methods << endl;
        }
    }
    
    
    // output1: num of 联通树
    cout << part_cnt << ' ' << num_methods;

    // output2: 每个联通树连接的非染色边的乘积

}
// 64 位输出请用 printf("%lld")