fork(1) download 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. const int N = 200000 + 5;
  5. const int LOG = 20;
  6.  
  7. int n, q;
  8. int val[N];
  9. vector<int> g[N];
  10. int up[N][LOG], depth[N];
  11. long long f[N];
  12.  
  13. int U[N], V[N], W[N];
  14. long long needv[N];
  15. long long ans[N];
  16.  
  17. vector<tuple<int,int,int>> gu[N];
  18.  
  19. struct BIT {
  20.     int n; vector<int> b;
  21.     void init(int _n){ n=_n; b.assign(n+1,0); }
  22.     void add(int i,int v){ for(; i<=n; i+=i&-i) b[i]+=v; }
  23.     int sum(int i){ int r=0; for(; i>0; i-=i&-i) r+=b[i]; return r; }
  24. } bit;
  25.  
  26. void dfs_init(int u,int p){
  27.     up[u][0]=p;
  28.     for(int j=1;j<LOG;j++) up[u][j]=up[up[u][j-1]][j-1];
  29.     for(int v:g[u]) if(v!=p){
  30.         depth[v]=depth[u]+1;
  31.         f[v]=f[u]+val[v];
  32.         dfs_init(v,u);
  33.     }
  34. }
  35.  
  36. int lca(int u,int v){
  37.     if(depth[u]<depth[v]) swap(u,v);
  38.     int k=depth[u]-depth[v];
  39.     for(int j=LOG-1;j>=0;j--) if(k>>j&1) u=up[u][j];
  40.     if(u==v) return u;
  41.     for(int j=LOG-1;j>=0;j--) if(up[u][j]!=up[v][j]){
  42.         u=up[u][j]; v=up[v][j];
  43.     }
  44.     return up[u][0];
  45. }
  46.  
  47. void dfs_ans(int u,int p){
  48.     static int M = 0;
  49.     bit.add((int)f[u],1);
  50.     for(auto [id, idxNeed, sign] : gu[u]){
  51.         ans[id] += 1LL*sign*(bit.sum(M) - bit.sum(idxNeed));
  52.     }
  53.     for(int v:g[u]) if(v!=p) dfs_ans(v,u);
  54.     bit.add((int)f[u],-1);
  55. }
  56.  
  57. int main(){
  58.     ios::sync_with_stdio(false);
  59.     cin.tie(nullptr);
  60.  
  61.     cin>>n>>q;
  62.     for(int i=1;i<=n;i++) cin>>val[i];
  63.     for(int i=1;i<n;i++){
  64.         int u,v; cin>>u>>v;
  65.         g[u].push_back(v);
  66.         g[v].push_back(u);
  67.     }
  68.  
  69.     f[1]=val[1];
  70.     dfs_init(1,0);
  71.  
  72.     vector<long long> vals;
  73.     vals.reserve(n+q);
  74.     for(int i=1;i<=n;i++) vals.push_back(f[i]);
  75.  
  76.     for(int i=1;i<=q;i++){
  77.         int u,v; cin>>u>>v;
  78.         int w=lca(u,v);
  79.         U[i]=u; V[i]=v; W[i]=w;
  80.         needv[i]=f[u]+f[v]-2*f[w]+val[w];
  81.         vals.push_back(needv[i]);
  82.     }
  83.  
  84.     sort(vals.begin(), vals.end());
  85.     vals.erase(unique(vals.begin(), vals.end()), vals.end());
  86.  
  87.     for(int i=1;i<=n;i++){
  88.         f[i] = lower_bound(vals.begin(), vals.end(), f[i]) - vals.begin() + 1;
  89.     }
  90.  
  91.     vector<int> idxNeed(q+1);
  92.     for(int i=1;i<=q;i++){
  93.         idxNeed[i] = upper_bound(vals.begin(), vals.end(), needv[i]) - vals.begin();
  94.         gu[U[i]].push_back({i, idxNeed[i], +1});
  95.         gu[V[i]].push_back({i, idxNeed[i], +1});
  96.         gu[W[i]].push_back({i, idxNeed[i], -2});
  97.         if((long long)(lower_bound(vals.begin(), vals.end(), f[W[i]]) - vals.begin() + 1) > idxNeed[i]) ans[i]++;
  98.     }
  99.  
  100.     bit.init((int)vals.size());
  101.     dfs_ans(1,0);
  102.  
  103.     for(int i=1;i<=q;i++) cout<<ans[i]<<"\n";
  104.     return 0;
  105. }
  106.  
Success #stdin #stdout 0.01s 22028KB
comments (?)
stdin 
8 5
-1 1 1 1 -1 1 1  -1
1 5
5 6
3 6
4 5
4 7
4 8
1 2
3 8
2 2
1 7
2 7
6 4
stdout 
-4
1
-2
-3
-2
https://ideone.com/IAnK11
4 4
-1
1
1
1
1 2
1 3
1 4
1 2
1 4
2 4
1 3

language:
C++14 (clang 8.0)
created:
5 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!