Distance on the tree（19南昌网络赛）

题目

https://nanti.jisuanke.com/t/38229

思路

1 主席树+LCA

#include <bits/stdc++.h>

using namespace std;
const int maxn = 200005;
struct edge
{
    int v,w;
    edge(int _v,int _w)
    {
        v = _v;
        w = _w;
    }
};
int n,m;
int SIZ = 1e9;
vector<edge> q[maxn];
int fa[maxn][20];
int deep[100005];
int tot,rt[100005];
struct TREE
{
    int l,r,x;
}tree[maxn*20];
void build(int &o,int l,int r,int data)
{
    tree[++tot] = tree[o];
    o = tot;
    tree[o].x++;
    if(l==r) return ;
    int mid = (l + r) / 2;
    if(mid >= data) build(tree[o].l,l,mid,data);
    else build(tree[o].r,mid+1,r,data);
}
void dfs(int u,int f,int d)
{
    deep[u] = d;
    fa[u][0] = f;
    int siz = q[u].size();
    for(int i = 0;i < siz;i++)
    {
        int v = q[u][i].v;
        int w = q[u][i].w;
        if(v == f) continue;
        rt[v] = rt[u];
        build(rt[v],0,SIZ,w);
        dfs(v,u,d+1);
    }
}
void init(int n)
{
    for(int j = 1;j < 19;j++)
    {
        for(int i = 1;i <= n;i++)
            fa[i][j] = fa[fa[i][j-1]][j-1];
    }
}
int lca(int u,int v)
{
    if(deep[u] < deep[v]) swap(u,v);
    int diff = deep[u] - deep[v];
    for(int i = 19;i >= 0;i--)
    {
        if(diff>>i&1)
        {
            u = fa[u][i];
        }
    }
    if(u == v) return u;
    for(int i = 19;i >= 0;i--)
    {
        if(fa[u][i] != fa[v][i])
        {
            u = fa[u][i];
            v = fa[v][i];
        }
    }
    return fa[u][0];
}
int Q(int s,int t,int l,int r,int L,int R)
{
    if(l >= L&&r <= R) return tree[t].x - tree[s].x;
    int ans = 0;
    int mid = (l + r) / 2;
    if(mid >= L) ans += Q(tree[s].l,tree[t].l,l,mid,L,R);
    if(mid < R) ans += Q(tree[s].r,tree[t].r,mid+1,r,L,R);
    return ans;
}
int query(int u,int v,int x)
{
    int uv = lca(u,v);
    return Q(rt[uv],rt[u],0,SIZ,0,x) + Q(rt[uv],rt[v],0,SIZ,0,x);
}
int main()
{
    scanf("%d%d",&n,&m);
    tot = 0;
    for(int i = 1;i < n;i++)
    {
        int u,v,w;
        scanf("%d%d%d",&u,&v,&w);
        q[u].push_back(edge(v,w));
        q[v].push_back(edge(u,w));
    }
    rt[0] = 0;
    tree[0].l = tree[0].r = tree[0].x = 0;
    dfs(1,0,1);
    init(n);
    while(m--)
    {
        int u,v,w;
        scanf("%d%d%d",&u,&v,&w);
        int ans = query(u,v,w);
        printf("%d\n",ans);
    }
    return 0;
}