ICPC 2019-2020 North-Western Russia Regional Contest E. Equidistant（分层）

#include <iostream>
#include <cstdio>
#include <algorithm>
#include <set>
#include <vector>
#include <queue>
#include <map>

using namespace std;

#define ll long long
#define pb push_back
#define fi first
#define se second

const int INF = 1e9;

vector<vector<int > > E;
vector<int > app;
vector<int > d;
vector<int > vis;
vector<int > cnt;

void bfs2(int n){
    fill(vis.begin(), vis.end(), 0);
    queue<int > que;
    for(int i = 1; i <= n; ++i){
        if(app[i]){
            que.push(i);
            vis[i] = 1;
        }
    }

    while(!que.empty()){
        int u = que.front();
        que.pop();
        for(auto v : E[u]){
            if(d[v] != d[u] + 1) continue; //只能下一层到上一层
            //printf(" u = %d v = %d\n", u, v);
            //vis[v] = 1;
            cnt[v] += cnt[u]; //汇聚
            if(!vis[v]){
                que.push(v);
                vis[v] = 1;
            }
        }
    }
}

void bfs1(int n ){
    queue<int > que;
    for(int i = 1 ;i <= n; ++i){
        if(app[i]) {
            que.push(i);
            d[i] = 1;
            vis[i] = 1;
        }
    }

    while(!que.empty()){
        int u = que.front();
        que.pop();
        for(auto v : E[u]){
            if(vis[v]) continue;
            vis[v] = 1;
            d[v] = d[u] + 1;
            que.push(v);
        }
    }
}

void solve(){
    int n, m;
    scanf("%d%d", &n, &m);
    E.resize(n + 10, vector<int >());
    d.resize(n + 10, 0);
    vis.resize(n + 10, 0);
    cnt.resize(n + 10, 0);
    app.resize(n + 10, 0); 
    int u, v;
    for(int i = 0; i < n - 1; ++i){
        scanf("%d%d", &u, &v);
        E[u].pb(v);
        E[v].pb(u);
    }
    //标记城市
    int city;
    for(int i = 0; i < m; ++i){
        scanf("%d", &city);
        app[city] = 1;
        cnt[city] = 1;
    }

    bfs1(n);//分层
    bfs2(n);//走层
    // for(int i = 1; i <= n; ++i) printf("u = %d d = %d\n", i, d[i]);
    // cout << endl;
    // for(int i = 1; i <= n; ++i) cout << cnt[i] << " ";
    // cout << endl;
    int inx = -1;
    for(int i = 1; i <= n; ++i){
        if(cnt[i] == m){
            inx = i; break;
        }
    }
    if(inx == -1) printf("NO\n");
    else printf("YSE\n%d\n", inx);
}

int main(){

    // ios::sync_with_stdio(false);
    // cin.tie(0); cout.tie(0);
    // freopen("C:\\Users\\admin\\Desktop\\input.txt", "r", stdin);
    // freopen("C:\\Users\\admin\\Desktop\\output.txt", "w", stdout);
    solve();
    //cout << "not error" << endl;
    return 0;
}