fork download
copy 
  1. // ~~ icebear ~~
  2. // Tag : Auxiliary Tree
  3. #include <bits/stdc++.h>
  4. using namespace std;
  5.  
  6. typedef long long ll;
  7. typedef pair<int, int> ii;
  8. typedef pair<int, ii> iii;
  9.  
  10. template<class T>
  11.     bool minimize(T &a, const T &b) {
  12.         if (a > b) return a = b, true;
  13.         return false;
  14.     }
  15.  
  16. template<class T>
  17.     bool maximize(T &a, const T &b) {
  18.         if (a < b) return a = b, true;
  19.         return false;
  20.     }
  21.  
  22. #define FOR(i,a,b) for(int i=(a); i<=(b); ++i)
  23. #define FORR(i,a,b) for(int i=(a); i>=(b); --i)
  24. #define REP(i, n) for(int i=0; i<(n); ++i)
  25. #define RED(i, n) for(int i=(n)-1; i>=0; --i)
  26. #define MASK(i) (1LL << (i))
  27. #define BIT(S, i) (((S) >> (i)) & 1)
  28. #define mp make_pair
  29. #define pb push_back
  30. #define fi first
  31. #define se second
  32. #define all(x) x.begin(), x.end()
  33. #define task "icebear"
  34. /*END OF TEMPLATE. ICEBEAR AND THE CAT WILL WIN VOI26 */
  35.  
  36. const int MOD = 1e9 + 7;
  37. const int inf = 1e9 + 27092008;
  38. const ll INF = 1e18 + 27092008;
  39. const int N = 2e5 + 5;
  40. int numNode, numQuery;
  41. vector<int> G[N];
  42. vector<ii> virG[N];
  43. int P[N][20], h[N], f[N], s[N], tin[N], tout[N], timer;
  44.  
  45. void dfs(int u) {
  46.     tin[u] = ++timer;
  47.     for(int v : G[u]) if (v != P[u][0]) {
  48.         h[v] = h[u] + 1;
  49.         P[v][0] = u;
  50.         dfs(v);
  51.     }
  52.     tout[u] = timer;
  53. }
  54.  
  55. int LCA(int u, int v) {
  56.     if (h[u] < h[v]) swap(u, v);
  57.     int s = h[u] - h[v];
  58.     RED(j, 20) if (BIT(s, j))
  59.         u = P[u][j];
  60.     if (u == v) return u;
  61.     RED(j, 20) if (P[u][j] != P[v][j]) {
  62.         u = P[u][j];
  63.         v = P[v][j];
  64.     }
  65.     return P[u][0];
  66. }
  67.  
  68. bool inSubtree(int u, int v) {
  69.     return tin[u] <= tin[v] && tout[v] <= tout[u];
  70. }
  71.  
  72. int ans, sumNode;
  73. void calcTree(int u) {
  74.     for(ii v : virG[u]) {
  75.         calcTree(v.fi);
  76.         f[u] = (f[u] + f[v.fi]) % MOD;
  77.         s[u] = (s[u] + s[v.fi]) % MOD;
  78.         ans = (ans + 1LL * (sumNode - f[v.fi] + MOD) * f[v.fi] % MOD * v.se % MOD) % MOD;
  79.     }
  80. }
  81.  
  82. int buildTree(int K, vector<int> &query) {
  83.     sort(all(query), [&](const int &x, const int &y){return tin[x] < tin[y];});
  84.     query.resize(unique(all(query)) - query.begin());
  85.     K = (int)query.size();
  86.  
  87.     sumNode = 0;
  88.     vector<int> initNode = query;
  89.     for(int &u : initNode) {
  90.         f[u] = u;
  91.         s[u] = 1;
  92.         (sumNode += u) %= MOD;
  93.     }
  94.  
  95.     FOR(i, 0, K - 2) query.pb(LCA(query[i], query[i + 1]));
  96.     sort(all(query), [&](const int &x, const int &y){return tin[x] < tin[y];});
  97.     query.resize(unique(all(query)) - query.begin());
  98.  
  99.     stack<int> st;
  100.     st.push(query[0]);
  101.     FOR(i, 1, (int)query.size() - 1) {
  102.         while(!inSubtree(st.top(), query[i])) st.pop();
  103.         virG[st.top()].emplace_back(query[i], h[query[i]] - h[st.top()]);
  104.         st.push(query[i]);
  105.     }
  106.  
  107.     ans = 0;
  108.     calcTree(query[0]);
  109.  
  110.     for(int &u : query) {
  111.         f[u] = s[u] = 0;
  112.         virG[u].clear();
  113.     }
  114.  
  115.     return ans;
  116. }
  117.  
  118. void init(void) {
  119.     cin >> numNode >> numQuery;
  120.     FOR(i, 2, numNode) {
  121.         int u, v;
  122.         cin >> u >> v;
  123.         G[u].pb(v);
  124.         G[v].pb(u);
  125.     }
  126. }
  127.  
  128. void process(void) {
  129.     dfs(1);
  130.     FOR(j, 1, 19) FOR(i, 1, numNode)
  131.         P[i][j] = P[P[i][j - 1]][j - 1];
  132.     while(numQuery--) {
  133.         int K; cin >> K;
  134.         vector<int> query(K);
  135.         for(int &u : query) cin >> u;
  136.         cout << buildTree(K, query) << '\n';
  137.     }
  138. }
  139.  
  140. int main() {
  141.     ios_base::sync_with_stdio(0);
  142.     cin.tie(0); cout.tie(0);
  143.     if (fopen(task".inp", "r")) {
  144.         freopen(task".inp", "r", stdin);
  145.         freopen(task".out", "w", stdout);
  146.     }
  147.     int tc = 1;
  148. //    cin >> tc;
  149.     while(tc--) {
  150.         init();
  151.         process();
  152.     }
  153.     return 0;
  154. }
  155.  
Success #stdin #stdout 0.01s 15600KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/SPugGb
language:
C++14 (gcc 8.3)
created:
1 day 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!