fork download
copy 
  1. #include <bits/stdc++.h>
  2.  
  3. using namespace std;
  4. const int N = 1e5 + 3;
  5. int head[N], heavy[N], depth[N], n, m, a[N], par[N], sz[N], t, pos[N];
  6. vector<int> g[N], mang[N];
  7. void DFS(int u){
  8.     sz[u] = 1;
  9.     int maxx = 0;
  10.     for(auto x: g[u]){
  11.         if(x == par[u]) continue;
  12.         depth[x] = depth[u] + 1;
  13.         par[x] = u;
  14.         DFS(x);
  15.         sz[u] += sz[x];
  16.         if(sz[x] > maxx){
  17.             maxx= sz[x];
  18.             heavy[u] = x;
  19.         }
  20.     }
  21.     return;
  22. }
  23. void buildhld(int u, int h){
  24.     head[u] = h;
  25.     pos[u] = ++t;
  26.     if(heavy[u]){
  27.         buildhld(heavy[u], h);
  28.     }
  29.     for(auto x: g[u]){
  30.         if(x == par[u] || x == heavy[u]) continue;
  31.         buildhld(x, x);
  32.     }
  33.     return;
  34. }
  35. bool solve(int u, int v, int val){
  36.     int res = 0;
  37.     while(head[u] != head[v]){
  38.         if(depth[head[u]] > depth[head[v]]) swap(u, v);
  39.         res += upper_bound(mang[val].begin(), mang[val].end(), pos[v])
  40.         - lower_bound(mang[val].begin(), mang[val].end(), pos[head[v]]);
  41.         v = par[head[v]];
  42.     }
  43.     if(depth[u] > depth[v]) swap(u, v);
  44.     res += upper_bound(mang[val].begin(), mang[val].end(), pos[v])
  45.         - lower_bound(mang[val].begin(), mang[val].end(), pos[u]);
  46.     return res > 0;
  47. }
  48. int main()
  49. {
  50.     ios_base::sync_with_stdio(false);
  51.     cin.tie(NULL);
  52.     cout.tie(NULL);
  53.     int u, v, i, w;
  54.     cin >> n >> m;
  55.     for(i = 1; i <= n; i++){
  56.         cin >> a[i];
  57.     }
  58.     for(i = 1; i < n; i++){
  59.         cin >> u >> v;
  60.         g[u].push_back(v);
  61.         g[v].push_back(u);
  62.     }
  63.     DFS(1);
  64.     buildhld(1, 1);
  65.     for(i = 1; i <= n; i++){
  66.         mang[a[i]].push_back(pos[i]);
  67.     }
  68.     for(i = 1; i <= n; i++) sort(mang[i].begin(), mang[i].end());
  69.     for(i = 1; i <= m; i++){
  70.         cin >> u >> v >> w;
  71.         cout << solve(u, v, w);
  72.     }
  73.     return 0;
  74. }
  75.  
Success #stdin #stdout 0.01s 9276KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
Standard output is empty
https://ideone.com/5VWHx0
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!