南昌网络赛 Distance on the tree（可持久化线段树）

题意：给一颗n个点的树，每条边都有一个权值， 每次询问u,v路径上权值小于k的边的数量

思路：感谢罗老板

代码：

#include <bits/stdc++.h>
using namespace std;

const int maxn = 1e5 + 7;
int nxt[maxn<<1],to[maxn<<1],head[maxn<<1],val[maxn<<1];
int cnt, tot;
void addedge(int u,int v,int w){
    nxt[++tot]=head[u];head[u]=tot;to[tot]=v;val[tot]=w;
}
struct node
{
  int l, r, sum;
}T[maxn * 40];
int rt[maxn * 40];
void update(int l, int r, int &x, int y, int pos)
{
  T[++cnt] = T[y]; T[cnt].sum++; x = cnt;
  if (l == r) return ;
  int mid = (l + r) >> 1;
  if (pos <= mid) update(l, mid, T[x].l, T[y].l, pos);
  else update(mid + 1, r, T[x].r, T[y].r, pos);
  T[x].sum = T[T[x].l].sum + T[T[x].r].sum;
}
int query(int i, int l, int r, int val)
{
  if (l == r) {
    return T[i].sum;
  }
  int mid = (l + r) >> 1;
  if(val <= mid) return query(T[i].l, l, mid, val);
  else return T[T[i].l].sum + query(T[i].r, mid + 1, r, val);
}
int fa[maxn][22], dep[maxn];
void dfs(int x, int pre, int deep)
{
  fa[x][0] = pre;
  dep[x] = deep;
  for (int i = head[x] ; i; i = nxt[i]) {
    int v = to[i];
    if (v == pre) continue;
    update(1, 1e9, rt[v], rt[x], val[i]);
    dfs(v, x, deep + 1);
  }
}
void init(int n)
{
  for (int i = 1; i <= 20; i++) {
    for (int j = 1; j <= n; j++) {
      int p = fa[j][i - 1];
      fa[j][i] = fa[p][i - 1];
    }
  }
}
int lca(int p, int q)
{
  if (dep[p] < dep[q])swap(p, q);
  for (int i = 20; i >= 0; i--) {
    if (dep[fa[p][i]] >= dep[q]) p = fa[p][i];  
  }
  if(p == q) return q;
  for (int i = 20; i >= 0; i--) {
    if (fa[p][i] != fa[q][i]) {
      p = fa[p][i]; q = fa[q][i];
    }
  }
  return fa[p][0];
}
int qu(int u, int v, int w)
{
  if(!w) return 0;
  int k = lca(u, v);
  int ans = query(rt[u],1, 1e9, w) + query(rt[v], 1, 1e9, w) - 2 * query(rt[k], 1, 1e9, w);
  return ans;
}
int main()
{
  int n, m;
  scanf("%d%d", &n, &m);
  for (int i = 2; i <= n; i++) {
    int u, v, w;
    scanf("%d%d%d", &u, &v, &w);
    addedge(u, v, w);
    addedge(v, u, w);
  }
  dfs(1, 1, 0);
  init(n);
  while(m --) {
    int u, v, w;
    scanf("%d%d%d", &u, &v, &w);
    printf("%d
", qu(u, v, w));
  }
  return 0;
}