fork(1) download
copy 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3. using ll = long long;
  4. const int MAXN = 500000 + 5;
  5. const int LOG = 20;
  6. const ll INFLL = (1LL<<60);
  7.  
  8. int N, Q;
  9. vector<int> tree[MAXN];
  10. int up[MAXN][LOG];
  11. int depth_[MAXN], tin[MAXN], tout[MAXN], timer_;
  12. int sz[MAXN];
  13.  
  14. void dfs_init(int u, int p){
  15.     tin[u] = ++timer_;
  16.     up[u][0] = p;
  17.     sz[u] = 1;
  18.     for(int k=1;k<LOG;k++) up[u][k] = up[ up[u][k-1] ][k-1];
  19.     for(int v: tree[u]) if(v!=p){
  20.         depth_[v] = depth_[u] + 1;
  21.         dfs_init(v, u);
  22.         sz[u] += sz[v];
  23.     }
  24.     tout[u] = ++timer_;
  25. }
  26.  
  27. bool is_anc(int u, int v){ // u ancestor of v
  28.     return tin[u] <= tin[v] && tout[v] <= tout[u];
  29. }
  30.  
  31. int lca(int u, int v){
  32.     if(is_anc(u,v)) return u;
  33.     if(is_anc(v,u)) return v;
  34.     for(int k=LOG-1;k>=0;k--){
  35.         if(!is_anc(up[u][k], v)) u = up[u][k];
  36.     }
  37.     return up[u][0];
  38. }
  39.  
  40. int dist_tree(int u, int v){
  41.     int w = lca(u,v);
  42.     return depth_[u] + depth_[v] - 2*depth_[w];
  43. }
  44.  
  45. // build virtual tree nodes from 'nodes_in' (contains specials + root)
  46. // returns nodes (unique, sorted by tin) and fills adj (indices 0..m-1)
  47. vector<int> build_virtual_tree(vector<int> nodes_in, vector<vector<int>> &adj){
  48.     // sort input by tin
  49.     sort(nodes_in.begin(), nodes_in.end(), [&](int a,int b){ return tin[a] < tin[b]; });
  50.     // add LCAs of adjacent nodes
  51.     int m0 = nodes_in.size();
  52.     for(int i=0;i<m0-1;i++){
  53.         nodes_in.push_back(lca(nodes_in[i], nodes_in[i+1]));
  54.     }
  55.     sort(nodes_in.begin(), nodes_in.end(), [&](int a,int b){ return tin[a] < tin[b]; });
  56.     nodes_in.erase(unique(nodes_in.begin(), nodes_in.end()), nodes_in.end());
  57.  
  58.     int m = (int)nodes_in.size();
  59.     adj.assign(m, {});
  60.     unordered_map<int,int> idx; idx.reserve(m*2);
  61.     for(int i=0;i<m;i++) idx[nodes_in[i]] = i;
  62.  
  63.     // stack to connect edges
  64.     vector<int> st;
  65.     st.push_back(nodes_in[0]);
  66.     for(int i=1;i<m;i++){
  67.         int u = nodes_in[i];
  68.         while(!st.empty() && !is_anc(st.back(), u)) st.pop_back();
  69.         if(st.empty()){
  70.             // defensive: shouldn't happen with LCAs added, but handle
  71.             st.push_back(u);
  72.             continue;
  73.         }
  74.         int p = st.back();
  75.         int pi = idx[p], ui = idx[u];
  76.         adj[pi].push_back(ui);
  77.         adj[ui].push_back(pi);
  78.         st.push_back(u);
  79.     }
  80.     return nodes_in;
  81. }
  82.  
  83. // Dijkstra on virtual tree (nodes: original ids vector, adj: adjacency by indices)
  84. // sources: vector of original ids (specials + root)
  85. vector<int> dijkstra_virtual(const vector<int> &nodes, const vector<vector<int>> &adj, const vector<int> &sources){
  86.     int m = (int)nodes.size();
  87.     unordered_map<int,int> idx; idx.reserve(m*2);
  88.     for(int i=0;i<m;i++) idx[nodes[i]] = i;
  89.  
  90.     vector<ll> dist(m, INFLL);
  91.     vector<int> owner(m, INT_MAX);
  92.  
  93.     // min-heap: (dist, owner, index)
  94.     using T = tuple<ll,int,int>;
  95.     priority_queue<T, vector<T>, greater<T>> pq;
  96.     for(int s: sources){
  97.         auto it = idx.find(s);
  98.         if(it == idx.end()) continue; // should not happen
  99.         int si = it->second;
  100.         if(dist[si] > 0 || (dist[si]==0 && owner[si] > s)){
  101.             dist[si] = 0;
  102.             owner[si] = s;
  103.             pq.emplace(0LL, s, si);
  104.         }
  105.     }
  106.  
  107.     while(!pq.empty()){
  108.         auto [d, own, u] = pq.top(); pq.pop();
  109.         if(d > dist[u]) continue;
  110.         if(d==dist[u] && own > owner[u]) continue;
  111.         for(int v : adj[u]){
  112.             ll nd = d + (ll)dist_tree(nodes[u], nodes[v]);
  113.             if(nd < dist[v] || (nd == dist[v] && own < owner[v])){
  114.                 dist[v] = nd;
  115.                 owner[v] = own;
  116.                 pq.emplace(nd, own, v);
  117.             }
  118.         }
  119.     }
  120.     return owner; // owner[i] = original id of special that claims nodes[i]
  121. }
  122.  
  123. int main(){
  124.     ios::sync_with_stdio(false);
  125.     cin.tie(nullptr);
  126.     cin >> N;
  127.     for(int i=0;i<N-1;i++){
  128.         int u,v; cin >> u >> v;
  129.         tree[u].push_back(v);
  130.         tree[v].push_back(u);
  131.     }
  132.     timer_ = 0;
  133.     depth_[1] = 0;
  134.     dfs_init(1, 1);
  135.  
  136.     cin >> Q;
  137.     vector<ll> ans(N+1, 0);
  138.     while(Q--){
  139.         int m; cin >> m;
  140.         vector<int> special(m);
  141.         for(int i=0;i<m;i++) cin >> special[i];
  142.         // save original order to print later
  143.         vector<int> original_special = special;
  144.  
  145.         // add root (1)
  146.         special.push_back(1);
  147.  
  148.         // build virtual tree
  149.         vector<vector<int>> adj;
  150.         vector<int> nodes = build_virtual_tree(special, adj);
  151.         int M = nodes.size();
  152.  
  153.         // run dijkstra on virtual tree
  154.         vector<int> owner = dijkstra_virtual(nodes, adj, special); // size M
  155.  
  156.         // root virtual tree at index 0 and compute children (directed)
  157.         vector<int> parent(M, -1);
  158.         vector<vector<int>> children(M);
  159.         // BFS/stack to set parent
  160.         vector<int> stk;
  161.         stk.reserve(M);
  162.         stk.push_back(0);
  163.         parent[0] = -2;
  164.         for(int i=0;i<(int)stk.size();i++){
  165.             int u = stk[i];
  166.             for(int v : adj[u]){
  167.                 if(parent[v] == -1){
  168.                     parent[v] = u;
  169.                     children[u].push_back(v);
  170.                     stk.push_back(v);
  171.                 }
  172.             }
  173.         }
  174.  
  175.         // compute segment sizes: seg[u] = sz[nodes[u]] - sum(sz[nodes[child]])
  176.         vector<ll> seg(M);
  177.         for(int i=0;i<M;i++) seg[i] = (ll)sz[nodes[i]];
  178.         for(int u=0; u<M; ++u){
  179.             for(int v : children[u]){
  180.                 seg[u] -= (ll)sz[nodes[v]];
  181.             }
  182.         }
  183.  
  184.         // add to answers: ans[ owner[i] ] += seg[i]
  185.         unordered_set<int> usedOwners;
  186.         usedOwners.reserve(M*2);
  187.         for(int i=0;i<M;i++){
  188.             int o = owner[i];
  189.             // safety: if owner[i] was not set (INT_MAX) - shouldn't happen because root included
  190.             if(o==INT_MAX) continue;
  191.             ans[o] += seg[i];
  192.             usedOwners.insert(o);
  193.         }
  194.  
  195.         // print answers in original order
  196.         for(int x : original_special){
  197.             cout << ans[x] << " ";
  198.         }
  199.         cout << "\n";
  200.  
  201.         // reset ans for used owners
  202.         for(int o : usedOwners) ans[o] = 0;
  203.     }
  204.     return 0;
  205. }
  206.  
Success #stdin #stdout 0.01s 23968KB
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 4
https://ideone.com/MJcKl8
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!