qCIh83 - Online C++0x Compiler & Debugging Tool

#include <bits/stdc++.h>
using namespace std;

const int maxn = 5e5 + 5;

vector<int> adj[maxn], g[maxn];
int par[maxn][20], h[maxn], sz[maxn];
int tin[maxn], tout[maxn], timeDfs = 0;

int ans[maxn];

array<int, 2> dp[maxn];

bool cmp(int u, int v){
    return tin[u] < tin[v];
}

bool inside(int u, int v){
    return (tin[u] <= tin[v] && tout[v] <= tout[u]);
}

void dfs(int u){
    tin[u] = ++timeDfs;
    sz[u] = 1;
    for(int v: adj[u]){
        if(v == par[u][0]) continue;
        par[v][0] = u;
        h[v] = h[u] + 1;

        for(int i = 1; (1 << i) <= h[v]; i++) par[v][i] = par[par[v][i - 1]][i - 1];

        dfs(v);
        sz[u] += sz[v];
    }
    tout[u] = timeDfs;
}

int up(int u, int d){
    assert(d >= 0);
    for(int i = 0; (1 << i) <= d; i++) if((d >> i) & 1) u = par[u][i];
    return u;
}

int lca(int u, int v){
    if(h[u] < h[v]) swap(u, v);

    u = up(u, h[u] - h[v]);

    if(u == v) return u;
    int lg = __lg(h[u]);

    for(int i = lg; i >= 0; i--){
        if(par[u][i] != par[v][i]){
            u = par[u][i];
            v = par[v][i];
        }
    }

    return par[u][0];
}


void solve(int u){
    for(int v: g[u]){
        solve(v);
        dp[u] = min(dp[u], {dp[v][0] + h[v] - h[u], dp[v][1]});
    }
    ans[dp[u][1]] += sz[u];
    for(int v: g[u]){
        if(dp[u][1] == dp[v][1]) ans[dp[u][1]] -= sz[v];
        else{
            int d = dp[u][0] + dp[v][0] + h[v] - h[u];
            //d = min dist u -> special vertice trong subtree u
                 //min dist v -> special vertice trong subtree v
                 //dist u, v
            if(d & 1) d /= 2;
            else{
                d /= 2;
                if(dp[u][1] < dp[v][1]) d--;
            }
            int mid = up(v, d - dp[v][0]);
            ans[dp[u][1]] -= sz[mid];
            ans[dp[v][1]] += sz[mid] - sz[v];
        }
    }
}

void solve(){
    int n; cin >> n;

    for(int i = 1; i <= n - 1; i++){
        int u, v; cin >> u >> v;
        adj[u].push_back(v);
        adj[v].push_back(u);
    }

    dfs(1);

    for(int i = 1; i <= n; i++) dp[i] = {(int)1e9, (int)1e9};
    int q; cin >> q;
    while(q--){
        int k; cin >> k;

        vector<int> x(k + 1), v;
        for(int i = 1; i <= k; i++){
            cin >> x[i];
            v.push_back(x[i]);
            dp[x[i]] = {0, x[i]};
        }

        sort(v.begin(), v.end(), cmp);
        int s = v.size();
        for(int i = 0; i < s - 1; i++) v.push_back(lca(v[i], v[i + 1]));

        v.push_back(1);
        sort(v.begin(), v.end(), cmp);
        v.resize(unique(v.begin(), v.end()) - v.begin());

        stack<int> st;
        for(int u: v){
            while(st.size() && !inside(st.top(), u)) st.pop();
            if(st.size()) g[st.top()].push_back(u);
            st.push(u);
        }

        solve(1);
        for(int i = 1; i <= k; i++) cout << ans[x[i]] << " ";
        cout << "\n";

        while(st.size()) st.pop();
        for(int u: v){
            g[u].clear();
            ans[u] = 0;
            dp[u] = {(int)1e9, (int)1e9};
        }
        v.clear();
    }
}


signed main() {
    ios_base::sync_with_stdio(false);
    cin.tie(0); cout.tie(0);

    solve();
}