2019南昌网络赛 J. Distance on the tree

我们对于每个节点建立到根节点的线段树，即我们建立一颗主席树来维护
对于每个查询，答案是query(1, x) + query(1, y) - 2 * query(1, lca(x, y))

#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef unsigned long long ull;
const int N = 1e5 + 5;
struct node//主席树
{
    int l;
    int r;
    int val;//存的是节点个数
}T[N * 50];
struct edge
{
    int to;
    int w;
};
int root[N], fa[N][21], deep[N], cnt, n, m;
vector <int> v;
vector <edge> G[N];
int getid(int x)
{
    return lower_bound(v.begin(), v.end(), x) - v.begin() + 1;
}
void update(int l, int r, int &x, int y, int pos)
{
    T[++ cnt] = T[y];
    T[cnt].val ++;
    x = cnt;
    if(l == r)
        return ;
    int mid = (l + r) >> 1;
    if(mid >= pos)
        update(l, mid, T[x].l, T[y].l, pos);
    else 
        update(mid + 1, r, T[x].r, T[y].r, pos);
}
int query(int l, int r, int x, int y, int k)
{
    if(r <= k)
        return T[y].val - T[x].val;
    int mid = (l + r) >> 1;
    if(k <= mid)
        return query(l, mid, T[x].l, T[y].l, k);
    else 
        return T[T[y].l].val - T[T[x].l].val + query(mid + 1, r, T[x].r, T[y].r, k);
}
void dfs(int now, int last)//遍历树，并建立主席树
{
    deep[now] = deep[last] + 1;
    fa[now][0] = last;
    for(int i = 1; (1 << i) <= deep[now]; i ++)
        fa[now][i] = fa[fa[now][i - 1]][i - 1];
    for(auto it : G[now])
    {
        int to = it.to, w = it.w;
        if(to == last)
            continue;
        update(1, n, root[to], root[now], getid(w));
        dfs(to, now);
    }
}
int lca(int x, int y)//树上倍增找最近公共祖先
{
    if(deep[x] > deep[y])
        swap(x, y);
    int d = deep[y] - deep[x];
    for(int i = 0; i <= 20; i ++)
        if((1 << i) & d)
            y = fa[y][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[y][0];
}
int main()
{
    scanf("%d%d", &n, &m);
    for(int i = 0; i < n - 1; i ++)
    {
        int u, to, w;
        scanf("%d%d%d", &u, &to, &w);
        G[u].push_back({to, w});
        G[to].push_back({u, w});
        v.push_back(w);
    }
    sort(v.begin(), v.end());
    v.erase(unique(v.begin(), v.end()), v.end());//离散化
    dfs(1, 0);
    while(m --)
    {
        int x, y, k;
        scanf("%d%d%d", &x, &y, &k);
        k = upper_bound(v.begin(), v.end(), k) - v.begin();
        if(!k)
            puts("0");//没有比k小的权值
        else 
            printf("%d\n", (query(1, n, root[1], root[x], k) + query(1, n, root[1], root[y], k) - 2 * query(1, n, root[1], root[lca(x, y)], k)));
    }
}