2019南昌网络赛 J Distance on the tree

主席树+lca

#include<bits/stdc++.h>
using namespace std;
const int maxn = 1e5 + 5;
int ver[maxn*2], edge[maxn*2], Next[maxn*2], head[maxn];
int a[maxn], T[maxn], dep[maxn], fa[maxn][25];
int tot, t, times, ans, cnt;
struct node1{ int l, r, sum; } tree[maxn*40];
inline int read()
{
    int s = 0, f = 1; char ch = getchar();
    while(ch < '0' || ch > '9'){ if(ch == '-') f = -1; ch = getchar(); }
    while(ch >= '0' && ch <= '9'){ s = s*10+ch-'0'; ch = getchar(); }
    return s*f;
}
void Add(int x,int y,int z)
{
    ver[++tot] = y; edge[tot] = z; Next[tot] = head[x]; head[x] = tot;
}
void init()
{
    T[0] = 0;
    tree[0].l = tree[0].r = tree[0].sum = 0;
}
void update(int &rt,int num,int l,int r)
{
    tree[++times] = tree[rt]; rt = times;
    tree[rt].sum++;
    if(l == r) return ;
    int mid = (l+r)>>1;
    if(num <= mid) update(tree[rt].l,num,l,mid);
    else update(tree[rt].r,num,mid+1,r);
}
void dfs(int u,int f,int z)
{
    T[u] = T[f];
    dep[u] = dep[f]+1;
    int p = lower_bound(a+1,a+cnt+1,z)-a;
    if(u != 1) update(T[u],p,1,maxn);
    for(int i = 0; i <= 19; i++) fa[u][i+1] = fa[fa[u][i]][i];
    for(int i = head[u],v; i; i = Next[i]){
        v = ver[i];
        if(v == f) continue;
        fa[v][0] = u;
        dfs(v,u,edge[i]);
    }
}
int lca(int x,int y)
{
    if(dep[x] < dep[y]) swap(x,y);
    for(int i = 20; i >= 0; i--){
        if(dep[fa[x][i]] >= dep[y]) x = fa[x][i];
        if(x == y) return x;
    }
    for(int i = 20; i >= 0; i--){
        if(fa[x][i] != fa[y][i]){
            x = fa[x][i];
            y = fa[y][i];
        }
    }
    return fa[x][0];
}
void query(int t1,int t2,int k,int l,int r)
{
    if(r <= k){
        ans += tree[t2].sum-tree[t1].sum;
        return ;
    }
    int mid = (l+r)>>1;
    if(k <= mid) query(tree[t1].l,tree[t2].l,k,l,mid);
    else{
        ans += tree[tree[t2].l].sum-tree[tree[t1].l].sum;
        query(tree[t1].r,tree[t2].r,k,mid+1,r);
    }
}
int main()
{
    int n, m, x, y, z;
    scanf("%d%d",&n,&m);
    for(int i = 1; i < n; i++){
        x = read(); y = read(); z = read();
        Add(x,y,z); Add(y,x,z);
        a[i] = z;
    }
    sort(a+1,a+n+1);
    cnt = unique(a+1,a+n+1)-a-1;
    init(); dfs(1,0,0);
    for(int i = 1,res; i <= m; i++){
        res = 0;
        x = read(); y = read(); z = read();
        if(!z) printf("0\n");
        else{
            int tmp = lca(x,y);
            z = upper_bound(a+1,a+cnt+1,z)-a-1;
            if(x != tmp) ans = 0, query(T[tmp],T[x],z,1,maxn), res += ans;
            if(y != tmp) ans = 0, query(T[tmp],T[y],z,1,maxn), res += ans;
            printf("%d\n",res);
        }
    }
    return 0;
}