2019 南昌网络赛 J. Distance on the tree (树链剖分 + 离线 + 线段树)

Description
DSM(Data Structure Master) once learned about tree when he was preparing for NOIP(National Olympiad in Informatics in Provinces) in Senior High School. So when in Data Structure Class in College, he is always absent-minded about what the teacher says.

The experienced and knowledgeable teacher had known about him even before the first class. However, she didn’t wish an informatics genius would destroy himself with idleness. After she knew that he was so interested in ACM(ACM International Collegiate Programming Contest), she finally made a plan to teach him to work hard in class, for knowledge is infinite.

This day, the teacher teaches about trees." A tree with nn nodes, can be defined as a graph with only one connected component and no cycle. So it has exactly $n-1$ edges…" DSM is nearly asleep until he is questioned by teacher. " I have known you are called Data Structure Master in Graph Theory, so here is a problem. “” A tree with nn nodes, which is numbered from $1$ to $n$. Edge between each two adjacent vertexes $u$ and $v$ has a value $w$, you’re asked to answer the number of edge whose value is no more than $k$ during the path between $u$ and $v$."" If you can’t solve the problem during the break, we will call you DaShaMao(Foolish Idiot) later on."

The problem seems quite easy for DSM. However, it can hardly be solved in a break. It’s such a disgrace if DSM can’t solve the problem. So during the break, he telephones you just for help. Can you save him for his dignity?

Input
In the first line there are two integers n,mn,m, represent the number of vertexes on the tree and queries
$(2\le n\le 10^5,1\le m\le 10^{5}2\le n\le 10^5,1\le m\le 10^5)$

The next $n-1$ lines, each line contains three integers $u,v,w$, indicates there is an undirected edge between nodes $u$ and $v$ with value $w$.
$(1\le u,v\le n,1\le w\le 10^9，1\le u,v\le n,1\le w\le 10^9)$

The next $m$ lines, each line contains three integers $u,v,k$ , be consistent with the problem given by the teacher above.
$(1\le u,v\le n,0\le k\le 10^9)(1\le u,v\le n,0\le k\le 10^9)$

Output
For each query, just print a single line contains the number of edges which meet the condition.

Main idea & Solution
给一颗有边权的树，有 m 次询问，每次询问u ~ v的简单路径上边权不大于k的边有多少条
将询问离线，按k值排序，树链剖分后用线段树维护。

Hint
因为这里是维护边的条数，这里采用将边权赋在其两端深度较大的端点上，也就是每个节点的值代表着它到其父节点的边权。
所以树链剖分时，需要特别注意答案的统计

Code

#include <bits/stdc++.h>
#define ls rt << 1
#define rs rt << 1 | 1
using namespace std;

typedef long long ll;
const int MX = 2e5 + 7;

int n,m;
int ecnt,head[MX];

struct Edge{
    int v,w,next;
}e[MX << 1];

void add(int u,int v,int w){
    e[++ecnt].v = v;e[ecnt].w = w;e[ecnt].next = head[u];head[u] = ecnt;
    e[++ecnt].v = u;e[ecnt].w = w;e[ecnt].next = head[v];head[v] = ecnt;
}

struct Segtree{
    int l,r,sum;
}t[MX << 3];

void push_up(int rt){
    t[rt].sum = t[ls].sum + t[rs].sum;
}

void build(int rt,int l,int r){
    t[rt].l = l, t[rt].r = r;
    t[rt].sum = 0;
    if(l == r) return ;
    int mid = (l + r) >> 1;
    build(ls,l,mid);
    build(rs,mid+1,r);
    push_up(rt);
}

void add(int rt,int pos){
    int l = t[rt].l, r = t[rt].r;
    if(l == r && l == pos){
        t[rt].sum = 1;
        return ;
    }
    int mid = (l + r) >> 1;
    if(pos <= mid) add(ls,pos);
    if(pos  > mid) add(rs,pos);
    push_up(rt);
}

int query(int rt,int L,int R){
    int l = t[rt].l, r = t[rt].r;
    if(L <= l && r <= R){
        return t[rt].sum;
    }
    int res = 0;
    int mid = (l + r) >> 1;
    if(L <= mid) res += query(ls,L,R);
    if(R  > mid) res += query(rs,L,R);
    return res; 
}

int siz[MX],son[MX],top[MX],fa[MX],id[MX],dep[MX],cnt;

void dfs1(int u,int f,int deep){
    siz[u] = 1, fa[u] = f, dep[u] = deep;
    int maxson = -1;
    for(int i = head[u];i;i = e[i].next){
        int v = e[i].v;
        if(v == f) continue;
        dfs1(v,u,deep+1);
        siz[u] += siz[v];
        if(siz[v] > maxson) son[u] = v, maxson = siz[v];
    }
}

void dfs2(int u,int topf){
    top[u] = topf, id[u] = ++cnt;
    if(!son[u]) return ;
    dfs2(son[u],topf);
    for(int i = head[u];i;i = e[i].next){
        int v = e[i].v;
        if(v == fa[u] || v == son[u]) continue;
        dfs2(v,v);
    }
}

int qRange(int u,int v){
    int res = 0;
    while(top[u] != top[v]){
        if(dep[top[u]] < dep[top[v]]) swap(u,v);
        res += query(1,id[top[u]],id[u]);
        u = fa[top[u]];
    }
    if(dep[u] > dep[v]) swap(u,v);
    res += query(1,id[u]+1,id[v]);//这里一定要记得+1
    return res;
}

struct que{
    int u,v,k,id;
    inline bool operator<(const que&it) const{
        return k < it.k;
    }
}q[MX],ed[MX<<1];

int ans[MX];
int main(){
    scanf("%d%d",&n,&m);
    for(int i = 1;i < n;++i){
        int u,v,w;scanf("%d%d%d",&u,&v,&w);add(u,v,w);
        ed[i].u = u, ed[i].v = v, ed[i].k = w;
    }sort(ed + 1,ed + n);
    for(int i = 1;i <= m;++i){
        scanf("%d%d%d",&q[i].u,&q[i].v,&q[i].k);
        q[i].id = i;
    }sort(q + 1,q + 1 + m);
    dfs1(1,1,0); dfs2(1,1);
    build(1,1,n);
    int j = 1;
    for(int i = 1;i <= m;++i){
        while(ed[j].k <= q[i].k && j < n){
            int u = ed[j].u, v = ed[j].v;
            if(dep[u] < dep[v]) swap(u,v);
            add(1,id[u]);
            j++;
        }
        ans[q[i].id] = qRange(q[i].u,q[i].v);
    }
    for(int i = 1;i <= m;++i) printf("%d\n", ans[i]);
    return 0;
}