2019年南昌邀请赛 - 网络赛J：Distance on the tree【树上主席树】

题目:

2019年南昌邀请赛 - 网络赛J：Distance on the tree

题意:

给定一颗带边权的树，Q 次询问，每次查询 u 到 v 的路径上边权 <= W 的边的数量

分析:

一眼树链剖分+主席树的模板题，但有个更好的写法：

和求树上两点间距离相似，只要计算出 u 到根的路径上 <= w 的边数 + v 到根路径上 <= w 的边数 - 2*LCA(u,v)到根的路径上 <= w 的边数即可；只需要建 n 颗主席树，每颗维护每个节点到根路径上边权的前缀和就可以了

代码：

#include <bits/stdc++.h>

using namespace std;
const int maxn = 1e5+25;
int n,nn,m,u,v,w,cnt,head[maxn],rt[maxn],c[maxn],num;
struct edge{
    int to,nxt,w;
}e[maxn<<1];
inline void add(int u,int v,int w){
    e[++cnt] = (edge){v,head[u],w};
    head[u] = cnt;
}
struct tree{
    int ls,rs,sum;
}T[maxn*80];
void build(int l,int r,int &x,int y,int v){
    T[++num] = T[y]; T[num].sum++; x = num;
    if(l == r) return ;
    int mid = (l+r) >> 1;
    if(v > mid) build(mid+1,r,T[x].rs,T[y].rs,v);
    else build(l,mid,T[x].ls,T[y].ls,v);
}
void query(int l,int r,int R,int x,int &sum){
    if(l > R) return ;
    if(r <= R){
        sum += T[x].sum;
        return ;
    }
    int mid = (l+r)>>1;
    query(l,mid,R,T[x].ls,sum);
    query(mid+1,r,R,T[x].rs,sum);
}
inline int getid(int x){
    int pos = lower_bound(c+1,c+n+1,x)-c;
    if(pos > n || c[pos] > x) pos--;
    return pos;
}
int d[maxn],dp[maxn][20];
void dfs(int x,int fa,int deep){
    d[x] = deep; dp[x][0] = fa;
    for(int i = head[x];i > 0;i = e[i].nxt){
        int v = e[i].to; if(v == fa) continue;
        build(1,n,rt[v],rt[x],getid(e[i].w));
        dfs(v,x,deep+1); 
    }
}
void init(){
    for(int i = 1;(1<<i) <= nn; ++i){
        for(int j = 1;j <= nn; ++j)
            dp[j][i] = dp[dp[j][i-1]][i-1];
    }
}
int LCA(int u,int v){
    if(d[u] < d[v]) swap(u,v);
    int h = d[u] - d[v];
    for(int i = 0;i < 20; ++i){
        if(h&(1<<i)) u = dp[u][i];
    }
    if(u == v) return u;
    for(int i = 19; ~i; --i){
        if(dp[u][i] != dp[v][i]) u=dp[u][i],v=dp[v][i];
    }
    return dp[u][0];
}
int main(){
    scanf("%d %d",&n,&m); nn = n;
    for(int i = 1;i < n; ++i){
        scanf("%d %d %d",&u,&v,&w);
        add(u,v,w); add(v,u,w); c[i] = w;
    }
    sort(c+1,c+n+1),n = unique(c+1,c+n+1) - c;
    dfs(1,0,1); init();
    while(m--){
        scanf("%d %d %d",&u,&v,&w);
        int val = getid(w),sum1 = 0,sum2 = 0;
        query(1,n,val,rt[v],sum1);
        query(1,n,val,rt[u],sum1);
        query(1,n,val,rt[LCA(u,v)],sum2);
        printf("%d\n",sum1-2*sum2);
    }
    return 0;
}