fork download
copy 
  1. #include <bits/stdc++.h>
  2.  
  3. using namespace std;
  4.  
  5. const int N = 200005;
  6. int n, q;
  7. int c[N];
  8. vector<int> adj[N];
  9.  
  10. int depth[N];
  11. int up[N][20];
  12. bool has_coin[N];
  13. int prr[N];
  14. int W = 0;
  15.  
  16. void dfs1(int u, int p) {
  17.     up[u][0] = p;
  18.     for (int i = 1; i < 20; i++) {
  19.         up[u][i] = up[up[u][i - 1]][i - 1];
  20.     }
  21.     has_coin[u] = c[u];
  22.     for (int v : adj[u]) {
  23.         if (v != p) {
  24.             depth[v] = depth[u] + 1;
  25.             dfs1(v, u);
  26.             has_coin[u] |= has_coin[v];
  27.         }
  28.     }
  29. }
  30.  
  31. void dfs2(int u, int p, int current_prr) {
  32.     if (has_coin[u]) {
  33.         prr[u] = u;
  34.     } else {
  35.         prr[u] = current_prr;
  36.     }
  37.     for (int v : adj[u]) {
  38.         if (v != p) {
  39.             dfs2(v, u, prr[u]);
  40.         }
  41.     }
  42. }
  43.  
  44. int get_lca(int u, int v) {
  45.     if (depth[u] < depth[v]) swap(u, v);
  46.     int diff = depth[u] - depth[v];
  47.     for (int i = 0; i < 20; i++) {
  48.         if ((diff >> i) & 1) {
  49.             u = up[u][i];
  50.         }
  51.     }
  52.     if (u == v) return u;
  53.     for (int i = 19; i >= 0; i--) {
  54.         if (up[u][i] != up[v][i]) {
  55.             u = up[u][i];
  56.             v = up[v][i];
  57.         }
  58.     }
  59.     return up[u][0];
  60. }
  61.  
  62. int main() {
  63.     ios_base::sync_with_stdio(0);
  64.     cin.tie(0);
  65.     cout.tie(0);
  66.  
  67.     if (!(cin >> n >> q)) return 0;
  68.  
  69.     int root = 1;
  70.     for (int i = 1; i <= n; i++) {
  71.         cin >> c[i];
  72.         if (c[i] == 1) root = i;
  73.     }
  74.  
  75.     for (int i = 1; i < n; i++) {
  76.         int u, v;
  77.         cin >> u >> v;
  78.         adj[u].push_back(v);
  79.         adj[v].push_back(u);
  80.     }
  81.  
  82.     depth[root] = 0;
  83.     dfs1(root, root);
  84.     dfs2(root, root, root);
  85.  
  86.     for (int i = 1; i <= n; i++) {
  87.         if (i != root && has_coin[i]) {
  88.             W++;
  89.         }
  90.     }
  91.  
  92.     for (int i = 0; i < q; i++) {
  93.         int a, b;
  94.         cin >> a >> b;
  95.         int lca = get_lca(a, b);
  96.         int dist = depth[a] + depth[b] - 2 * depth[lca];
  97.  
  98.         if (has_coin[lca]) {
  99.             long long ans = 2LL * W + 2LL * (depth[a] - depth[prr[a]]) + 2LL * (depth[b] - depth[prr[b]]) - dist;
  100.             cout << ans << "\n";
  101.         } else {
  102.             long long ans = 2LL * W + dist + 2LL * (depth[lca] - depth[prr[a]]);
  103.             cout << ans << "\n";
  104.         }
  105.     }
  106.  
  107.     return 0;
  108. }
  109.  
Success #stdin #stdout 0.01s 8784KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/TqE8bD
language:
C++14 (clang 8.0)
created:
7 hours 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!