fork download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3. using ll = long long;
  4. const ll INF = (1LL<<60);
  5. const int MAXN = 500000 + 5;
  6. const int LOG = 20;
  7.  
  8. int N;
  9. vector<int> tree[MAXN];
  10. int up[MAXN][LOG];
  11. int tin[MAXN], tout[MAXN], timerdfs;
  12. int depthArr[MAXN];
  13. int subSize[MAXN];
  14.  
  15. void dfs_prep(int u, int p){
  16.     tin[u] = ++timerdfs;
  17.     up[u][0] = (p == -1 ? u : p);
  18.     for(int k=1;k<LOG;k++) up[u][k] = up[ up[u][k-1] ][k-1];
  19.     subSize[u] = 1;
  20.     for(int v: tree[u]) if(v != p){
  21.         depthArr[v] = depthArr[u] + 1;
  22.         dfs_prep(v, u);
  23.         subSize[u] += subSize[v];
  24.     }
  25.     tout[u] = ++timerdfs;
  26. }
  27. inline bool is_anc(int a, int b){ // a ancestor of b
  28.     return tin[a] <= tin[b] && tout[b] <= tout[a];
  29. }
  30. int lca(int a, int b){
  31.     if(is_anc(a,b)) return a;
  32.     if(is_anc(b,a)) return b;
  33.     for(int k=LOG-1;k>=0;k--){
  34.         if(!is_anc(up[a][k], b)) a = up[a][k];
  35.     }
  36.     return up[a][0];
  37. }
  38. inline int dist_tree(int a, int b){
  39.     int L = lca(a,b);
  40.     return depthArr[a] + depthArr[b] - 2*depthArr[L];
  41. }
  42.  
  43. int idxArr[MAXN]; // mapping node -> index in virtual nodes (reset only for touched)
  44. int main(){
  45.     ios::sync_with_stdio(false);
  46.     cin.tie(nullptr);
  47.     cin >> N;
  48.     for(int i=1;i<=N;i++){
  49.         tree[i].clear();
  50.         idxArr[i] = -1;
  51.     }
  52.     for(int i=0;i<N-1;i++){
  53.         int u,v; cin >> u >> v;
  54.         tree[u].push_back(v);
  55.         tree[v].push_back(u);
  56.     }
  57.     timerdfs = 0;
  58.     depthArr[1] = 0;
  59.     dfs_prep(1, -1);
  60.  
  61.     int Q; cin >> Q;
  62.     vector<ll> globalAns(N+1, 0);
  63.  
  64.     while(Q--){
  65.         int m; cin >> m;
  66.         vector<int> specials(m);
  67.         for(int i=0;i<m;i++) cin >> specials[i];
  68.  
  69.         // build virtual node list
  70.         vector<int> nodes = specials;
  71.         sort(nodes.begin(), nodes.end(), [&](int a,int b){ return tin[a] < tin[b]; });
  72.         int orig = nodes.size();
  73.         for(int i=0;i<orig-1;i++){
  74.             nodes.push_back(lca(nodes[i], nodes[i+1]));
  75.         }
  76.         sort(nodes.begin(), nodes.end(), [&](int a,int b){ return tin[a] < tin[b]; });
  77.         nodes.erase(unique(nodes.begin(), nodes.end()), nodes.end());
  78.         int K = nodes.size();
  79.  
  80.         // map nodes -> indices (and remember touched to reset later)
  81.         vector<int> touched = nodes;
  82.         for(int i=0;i<K;i++) idxArr[nodes[i]] = i;
  83.  
  84.         // build virtual tree adjacency (by indices), weighted edges = dist between original nodes
  85.         vector<vector<pair<int,int>>> adj(K);
  86.         vector<int> st;
  87.         st.push_back(nodes[0]);
  88.         for(int i=1;i<K;i++){
  89.             int cur = nodes[i];
  90.             while((int)st.size() > 1 && !is_anc(st.back(), cur)){
  91.                 int v = st.back(); st.pop_back();
  92.                 int p = st.back();
  93.                 int vi = idxArr[v], pi = idxArr[p];
  94.                 int w = dist_tree(p, v);
  95.                 adj[pi].push_back({vi, w});
  96.                 adj[vi].push_back({pi, w});
  97.             }
  98.             if(!is_anc(st.back(), cur)){
  99.                 // now st.size()==1 and still not ancestor: connect and continue
  100.                 int v = st.back(); st.pop_back();
  101.                 if(!st.empty()){
  102.                     int p = st.back();
  103.                     int vi = idxArr[v], pi = idxArr[p];
  104.                     int w = dist_tree(p, v);
  105.                     adj[pi].push_back({vi, w});
  106.                     adj[vi].push_back({pi, w});
  107.                 }
  108.                 st.push_back(v);
  109.             }
  110.             st.push_back(cur);
  111.         }
  112.         while(st.size() > 1){
  113.             int v = st.back(); st.pop_back();
  114.             int p = st.back();
  115.             int vi = idxArr[v], pi = idxArr[p];
  116.             int w = dist_tree(p, v);
  117.             adj[pi].push_back({vi, w});
  118.             adj[vi].push_back({pi, w});
  119.         }
  120.  
  121.         // multi-source Dijkstra on virtual tree
  122.         using T = tuple<ll,int,int>; // dist, owner(original id), index
  123.         priority_queue<T, vector<T>, greater<T>> pq;
  124.         vector<ll> dist(K, INF);
  125.         vector<int> owner(K, INT_MAX);
  126.         for(int s : specials){
  127.             int si = idxArr[s];
  128.             if(si < 0) continue;
  129.             if(dist[si] > 0 || (dist[si]==0 && s < owner[si])){
  130.                 dist[si] = 0;
  131.                 owner[si] = s;
  132.                 pq.emplace(0LL, s, si);
  133.             }
  134.         }
  135.         while(!pq.empty()){
  136.             auto [dcur, own, u] = pq.top(); pq.pop();
  137.             if(dcur != dist[u] || own != owner[u]) continue;
  138.             for(auto &e : adj[u]){
  139.                 int v = e.first;
  140.                 int w = e.second;
  141.                 ll nd = dcur + (ll)w;
  142.                 if(nd < dist[v] || (nd == dist[v] && own < owner[v])){
  143.                     dist[v] = nd;
  144.                     owner[v] = own;
  145.                     pq.emplace(nd, own, v);
  146.                 }
  147.             }
  148.         }
  149.  
  150.         // orient virtual tree as rooted at index 0 to compute children
  151.         vector<int> parent(K, -1);
  152.         vector<vector<int>> children(K);
  153.         parent[0] = -2;
  154.         vector<int> bfs;
  155.         bfs.push_back(0);
  156.         for(size_t i=0;i<bfs.size();i++){
  157.             int u = bfs[i];
  158.             for(auto &e: adj[u]){
  159.                 int v = e.first;
  160.                 if(parent[v] == -1){
  161.                     parent[v] = u;
  162.                     children[u].push_back(v);
  163.                     bfs.push_back(v);
  164.                 }
  165.             }
  166.         }
  167.  
  168.         // compute segment sizes and add to answers
  169.         vector<ll> seg(K);
  170.         for(int i=0;i<K;i++) seg[i] = subSize[ nodes[i] ];
  171.         for(int u=0; u<K; ++u){
  172.             for(int v : children[u]){
  173.                 seg[u] -= subSize[ nodes[v] ];
  174.             }
  175.         }
  176.  
  177.         // collect used owners to reset later
  178.         vector<int> usedOwners;
  179.         usedOwners.reserve(K);
  180.         for(int i=0;i<K;i++){
  181.             int o = owner[i];
  182.             if(o == INT_MAX) continue; // shouldn't happen
  183.             globalAns[o] += seg[i];
  184.             usedOwners.push_back(o);
  185.         }
  186.  
  187.         // print answers in input order for this query
  188.         for(int s : specials){
  189.             cout << globalAns[s] << " ";
  190.         }
  191.         cout << "\n";
  192.  
  193.         // reset globalAns for owners used (unique)
  194.         sort(usedOwners.begin(), usedOwners.end());
  195.         usedOwners.erase(unique(usedOwners.begin(), usedOwners.end()), usedOwners.end());
  196.         for(int o : usedOwners) globalAns[o] = 0;
  197.  
  198.         // reset idxArr for touched nodes
  199.         for(int u : touched) idxArr[u] = -1;
  200.     }
  201.  
  202.     return 0;
  203. }
  204.  
Success #stdin #stdout 0.01s 25984KB
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 4 4
https://ideone.com/CTuGjl
language:
C++14 (clang 8.0)
created:
25 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!