Distance on the tree（2019南昌邀请赛网络赛 J 题）

题目链接

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

题目大意

一棵树，问 u 到 v 上权值小于 k 的边的条数。

解题思路

LCA+主席树
因为查询1e5，所以用主席树预处理好后，lca找到最近公共祖先 uv，主席树查找 u 到 uv 的小于 k 的条数加 uv到 v 的小于 k 的条数。

代码

#include <iostream>
#include <bits/stdc++.h>
#define mid ((l+r)>>1)
using namespace std;
const int N=2e5+100;
vector<int>nxt[N],len[N],k;
int dep[N],fa[N][20];
int rt[N];
int tot,SIZ=5e8;
int n,m;
struct node
{
    int l,r,x;
}tree[N*20];

void updata(int &o,int l,int r,int data)
{
    tree[++tot]=tree[o];
    o=tot;
    tree[o].x++;
    if(l==r)
        return;
    if(mid>=data)
        updata(tree[o].l,l,mid,data);
    else updata(tree[o].r,mid+1,r,data);
}

int query(int lo,int ro,int l,int r,int h,int t)
{
    if(l>=h&&r<=t)
        return tree[ro].x-tree[lo].x;
    int ans=0;
    if(mid>=h)ans+=query(tree[lo].l,tree[ro].l,l,mid,h,t);
    if(mid<t)ans+=query(tree[lo].r,tree[ro].r,mid+1,r,h,t);
    return ans;
}
//lca 模板 start
void dfs(int u,int f,int d)
{
    dep[u]=d;
    fa[u][0]=f;
    int siz=nxt[u].size();
    for(int i=0;i<siz;i++)
    {
        int v=nxt[u][i];
        int w=len[u][i];
        if(v==f)continue;
        rt[v]=rt[u];
        updata(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(dep[u]<dep[v])
        swap(u,v);
    int dif=dep[u]-dep[v];
    for(int i=19;i>=0;i--)
    {
        if(dif>>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];
}
//lca 模板 end
int Query(int u,int v,int x)
{
    int uv=lca(u,v);
    return query(rt[uv],rt[u],0,SIZ,0,x)+query(rt[uv],rt[v],0,SIZ,0,x);
}

int main()
{
    scanf("%d%d",&n,&m);
    //建图 start
    for(int i=1;i<n;i++)
    {
        int u,v,w;
        scanf("%d%d%d",&u,&v,&w);
        nxt[u].push_back(v);
        nxt[v].push_back(u);
        len[u].push_back(w);
        len[v].push_back(w);
    }
    //建图 end
    //lca初始化 start
    dfs(1,0,1);
    init(n);
    //lca初始化 end
    rt[0]=0;
    tree[0].l=tree[0].r=tree[0].x=0;
    for(int i=1;i<=m;i++)
    {
        int u,v,x;
        scanf("%d%d%d",&u,&v,&x);
        printf("%d\n",Query(u,v,x));
    }
    return 0;
}