2019年ICPC南昌网络赛 J. Distance on the tree 树链剖分+主席树

#include<iostream>
#include<algorithm>
#include<stdio.h>
#include<string.h>
using namespace std;
const int maxx = 2e5+10;
struct node{
  int l,r,cnt;
}tree[maxx*40];
int head[maxx],rk[maxx],siz[maxx],top[maxx],son[maxx],d[maxx],fa[maxx],id[maxx],rt[maxx];
int dis[maxx];
struct Edge{
  int Next,to,w;
}edge[maxx<<1];
int root[maxx];
int tot=0,cnt=0,order;
void inserts(int l,int r,int pre,int &now,int pos){
    now=++cnt;
    tree[now]=tree[pre];
    tree[now].cnt++;
    if (l==r){
        return ;
    }
    int mid=(l+r)>>1;
    if (pos<=mid){
        inserts(l,mid,tree[pre].l,tree[now].l,pos);
    }else {
        inserts(mid+1,r,tree[pre].r,tree[now].r,pos);
    }
}
int query(int L,int R,int l,int r,int w){
    if (l==r){
        return tree[R].cnt-tree[L].cnt;
    }
    int mid=(l+r)>>1;
    if (w<=mid){
        return query(tree[L].l,tree[R].l,l,mid,w);
    }else {
        return tree[tree[R].l].cnt-tree[tree[L].l].cnt+query(tree[L].r,tree[R].r,mid+1,r,w);
    }
}
void add(int x,int y,int z){
  edge[++tot].to=y;
  edge[tot].w=z;
  edge[tot].Next=head[x];
  head[x]=tot;
}
void dfs1(int u,int f,int depth){
    d[u]=depth;
    fa[u]=f;
    siz[u]=1;
    for (int i=head[u];i;i=edge[i].Next){
        int v=edge[i].to;
        if (v==f)continue;
        dfs1(v,u,depth+1);
        dis[v]=edge[i].w;
        siz[u]+=siz[v];
        if(siz[v]>siz[son[u]])
            son[u]=v;
    }
}
void dfs2(int u,int t){
   top[u]=t;
   id[u]=++order;
   rk[order]=u;
  // cout<<dis[u]<<endl;
   inserts(0,1e9,root[order-1],root[order],dis[u]);
   if(!son[u])return;
   dfs2(son[u],t);
   for (int i=head[u];i;i=edge[i].Next)
   {
       int v=edge[i].to;
       if(v!=son[u] && v!=fa[u])
        dfs2(v,v);
   }
}
int query_line(int a,int b,int c){
  int ans=0;
  while(top[a]!=top[b]){
    if (d[top[a]]>d[top[b]])swap(a,b);
    ans+=query(root[id[top[b]]-1],root[id[b]],0,1e9,c);
    b=fa[top[b]];
  }
  if (d[a]>d[b])swap(a,b);
  ans+=query(root[id[a]],root[id[b]],0,1e9,c);
  return ans;
}
int main(){
  int w,op,n,uu,vv;
  scanf("%d%d",&n,&op);
  tot=0;
  cnt=0;
  order=0;
  for (int i=1;i<n;i++){
     scanf("%d%d%d",&uu,&vv,&w);
     add(uu,vv,w);
     add(vv,uu,w);
  }
  dfs1(1,0,1);
  dfs2(1,1);
//  cout<<endl;
//  for (int i=1;i<=n;i++){
//    cout<<dis[i]<<"--";
//  }
//  cout<<endl;
//  cout<<endl;
//  for (int i=1;i<=n;i++){
//    cout<<"--"<<id[i]<<endl;
//  }
  while(op--){
    scanf("%d%d%d",&uu,&vv,&w);
    printf("%d\n",w?query_line(uu,vv,w):0);
  }
  return 0;
}