fork(1) download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. const int MAXN = 500000;
  5. int N, Q;
  6. vector<int> tree[MAXN+1];
  7. int up[MAXN+1][20], depth_[MAXN+1], tin[MAXN+1], tout[MAXN+1], timer_;
  8. int sz[MAXN+1];
  9.  
  10. void dfs_init(int u, int p) {
  11.     tin[u] = ++timer_;
  12.     up[u][0] = p;
  13.     sz[u] = 1;
  14.     for (int k = 1; k < 20; k++)
  15.         up[u][k] = up[ up[u][k-1] ][k-1];
  16.     for (int v : tree[u]) if (v != p) {
  17.         depth_[v] = depth_[u] + 1;
  18.         dfs_init(v, u);
  19.         sz[u] += sz[v];
  20.     }
  21.     tout[u] = ++timer_;
  22. }
  23.  
  24. bool is_anc(int u, int v) { // u ancestor of v?
  25.     return tin[u] <= tin[v] && tout[v] <= tout[u];
  26. }
  27.  
  28. int lca(int u, int v) {
  29.     if (is_anc(u, v)) return u;
  30.     if (is_anc(v, u)) return v;
  31.     for (int k = 19; k >= 0; k--)
  32.         if (!is_anc(up[u][k], v))
  33.             u = up[u][k];
  34.     return up[u][0];
  35. }
  36.  
  37. int dist_tree(int u, int v) {
  38.     int w = lca(u, v);
  39.     return depth_[u] + depth_[v] - 2*depth_[w];
  40. }
  41.  
  42. vector<int> build_virtual_tree(vector<int> nodes, vector<vector<int>> &adj) {
  43.     sort(nodes.begin(), nodes.end(), [&](int a, int b){
  44.         return tin[a] < tin[b];
  45.     });
  46.     int m = nodes.size();
  47.     // Thêm các LCA vào
  48.     for (int i = 0; i < m-1; i++)
  49.         nodes.push_back(lca(nodes[i], nodes[i+1]));
  50.     sort(nodes.begin(), nodes.end(), [&](int a, int b){
  51.         return tin[a] < tin[b];
  52.     });
  53.     nodes.erase(unique(nodes.begin(), nodes.end()), nodes.end());
  54.  
  55.     // Xây cây ảo
  56.     adj.assign(nodes.size(), {});
  57.     vector<int> st;
  58.     vector<int> idx(MAXN+1, -1);
  59.     for (int i = 0; i < (int)nodes.size(); i++) idx[nodes[i]] = i;
  60.  
  61.     st.push_back(nodes[0]);
  62.     for (int i = 1; i < (int)nodes.size(); i++) {
  63.         int u = nodes[i];
  64.         while (!is_anc(st.back(), u)) st.pop_back();
  65.         int p = st.back();
  66.         adj[idx[p]].push_back(idx[u]);
  67.         adj[idx[u]].push_back(idx[p]);
  68.         st.push_back(u);
  69.     }
  70.     return nodes;
  71. }
  72.  
  73. struct Info {
  74.     int dist, id;
  75. };
  76.  
  77. vector<int> bfs_virtual(const vector<int> &nodes, const vector<vector<int>> &adj, const vector<int> &special) {
  78.     int m = nodes.size();
  79.     vector<Info> info(m, {INT_MAX, INT_MAX});
  80.     vector<int> idx(MAXN+1, -1);
  81.     for (int i = 0; i < m; i++) idx[nodes[i]] = i;
  82.  
  83.     queue<int> q;
  84.     for (int sp : special) {
  85.         int id = idx[sp];
  86.         info[id] = {0, sp};
  87.         q.push(id);
  88.     }
  89.     while (!q.empty()) {
  90.         int u = q.front(); q.pop();
  91.         for (int v : adj[u]) {
  92.             int d = info[u].dist + dist_tree(nodes[u], nodes[v]);
  93.             if (d < info[v].dist || (d == info[v].dist && info[u].id < info[v].id)) {
  94.                 info[v] = {d, info[u].id};
  95.                 q.push(v);
  96.             }
  97.         }
  98.     }
  99.  
  100.     vector<int> id_of(m);
  101.     for (int i = 0; i < m; i++) id_of[i] = info[i].id;
  102.     return id_of;
  103. }
  104.  
  105. long long ans[MAXN+1];
  106. int dp[MAXN+1];
  107.  
  108. int dfs_dp(int u, int p, const vector<int> &nodes, const vector<vector<int>> &adj, const vector<int> &id_of) {
  109.     int orig_u = nodes[u];
  110.     dp[u] = sz[orig_u];
  111.     for (int v : adj[u]) if (v != p) {
  112.         int sub = dfs_dp(v, u, nodes, adj, id_of);
  113.         if (id_of[v] == id_of[u]) dp[u] += sub;
  114.     }
  115.     ans[id_of[u]] += dp[u];
  116.     return dp[u];
  117. }
  118.  
  119. int main(){
  120.     ios::sync_with_stdio(false);
  121.     cin.tie(nullptr);
  122.     cin >> N;
  123.     for (int i = 1; i <= N-1; i++) {
  124.         int u, v; cin >> u >> v;
  125.         tree[u].push_back(v);
  126.         tree[v].push_back(u);
  127.     }
  128.     depth_[1] = 0;
  129.     dfs_init(1, 1);
  130.  
  131.     cin >> Q;
  132.     while (Q--) {
  133.         int m; cin >> m;
  134.         vector<int> special(m);
  135.         for (int i = 0; i < m; i++) cin >> special[i];
  136.         special.push_back(1); // thêm root
  137.  
  138.         vector<vector<int>> adj;
  139.         vector<int> nodes = build_virtual_tree(special, adj);
  140.  
  141.         vector<int> id_of = bfs_virtual(nodes, adj, special);
  142.  
  143.         dfs_dp(0, -1, nodes, adj, id_of);
  144.  
  145.         for (int x : special) if (x != 1) {
  146.             cout << ans[x] << " ";
  147.             ans[x] = 0; // reset
  148.         }
  149.         cout << "\n";
  150.     }
  151. }
  152.  
Success #stdin #stdout 0.01s 27908KB
comments (?)
stdin
copy 
10
2 1
3 2
4 3
5 4
6 1
7 3
8 3
9 4
10 1
1
4
8 7 10 3
stdout
copy 
1 1 1 6
https://ideone.com/Ilb1nx
language:
C++14 (clang 8.0)
created:
2 days ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!