[湖南集训]谈笑风生 线段树合并
给一棵树\(n\)个节点，\(q\)次询问，每次给定\(p,k\)，问有多少三元组\((p,b,c)\)满足\(p,b\)均为\(c\)的父亲，\(p,b\)在树上的距离不超过\(k\)
\(n,q\le 10^5\)
很有意思的一道题。
两种情况讨论：
- \(b\)在\(a\)的上面，我们发现\(b\)有\(min(dep[a]-1, k)\)个可选位置，而\(a\)的子树内所有节点（除了节点\(a\)本身）均可作为\(c\)，所以共有\(min(dep[a]-1, k)\times(sz[a]-1) \)个三元组满足。
- 对于\(b\)在\(a\)的下面，我们可以每个节点都维护一颗下标为深度的权值线段树，因为节点\(u\)作为\(b\)的同时，其子树除\(u\)外均可作为\(c\)，所以对于深度\(dep[u]\)的贡献为\(sz[u]-1\)，然后在\(dfs\)回溯时同时合并子树线段树，每次答案即为深度区间\([dep[a]+1, dep[a]+k]\)的区间和。
最后答案即为两种情况之和
注意空间限制，需要动态开点。
#include <cstdio>
#include <algorithm>
#define MAXN 300003
#define MAXM 300003*30
#define ll long long
using namespace std;
int head[MAXN],nxt[MAXN*2],vv[MAXN*2],tot;
inline void add_edge(int u, int v){
vv[++tot]=v;
nxt[tot]=head[u];
head[u]=tot;
}
int cnt;
ll tre[MAXM*2];
int sl[MAXM*2],sr[MAXM*2];
void change(int &x, int l, int r, int pos, int val){
if(x==0) x=++cnt;
tre[x]+=val;
if(l==r) return;
int mid=(l+r)>>1;
if(pos<=mid) change(sl[x], l, mid, pos, val);
else change(sr[x], mid+1, r, pos, val);
}
ll query(int x, int l, int r, int ql, int qr){
if(x==0) return 0;
if(ql<=l&&r<=qr) return tre[x];
int mid=(l+r)>>1;
ll res=0;
if(ql<=mid) res+=query(sl[x], l, mid, ql, qr);
if(mid<qr) res+=query(sr[x], mid+1, r, ql, qr);
return res;
}
int merge(int a, int b, int l, int r){
if(a==0||b==0) return a+b;
int mid=(l+r)>>1;
int x=++cnt;
tre[x]=tre[a]+tre[b];
sl[x]=merge(sl[a], sl[b], l, mid);
sr[x]=merge(sr[a], sr[b], mid+1, r);
return x;
}
int n,q;
int sz[MAXN],rot[MAXN],dep[MAXN];
void dfs(int u, int fa){
sz[u]=1;
dep[u]=dep[fa]+1;
for(int i=head[u];i;i=nxt[i]){
int v=vv[i];
if(v==fa) continue;
dfs(v, u);
sz[u]+=sz[v];
}
change(rot[u], 1, n, dep[u], sz[u]-1);
rot[fa]=merge(rot[fa], rot[u], 1, n);
}
int main(){
scanf("%d %d", &n, &q);
for(int i=1;i<n;++i){
int a,b;scanf("%d %d", &a, &b);
add_edge(a, b);add_edge(b, a);
}
dfs(1, 0);
while(q--){
int p,k;scanf("%d %d", &p, &k);
printf("%lld\n", query(rot[p], 1, n, dep[p]+1, dep[p]+k)+(ll)(sz[p]-1)*min(k, dep[p]-1));
}
return 0;
}