P3899-[湖南集训]谈笑风生【主席树】

正题

题目链接:https://www.luogu.com.cn/problem/P3899

题目大意

给出$n$个点的一棵有根树，每次询问一个$(p,k)$，求有多少个点对$(b,c)$满足

1. $p$和$b$是$c$的祖先
2. $b$与$p$的距离不超过$k$
   $$蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤蛤$$

解题思路

首先如果$b$在$a$上方，那么点$b$的个数可以用深度来求，而$c$的数量就是$a$的子树大小$?1$

如果$b$在$a$的下方，$c$的数量就是$b$的子树大小$?1$，也就是对于每个$b$它的贡献是它的子树大小$?1$，那么就求我们在$a$的子树中与$a$距离不超过$k$的点的权值和即可。

这个可以用$dfs$序和主席树维护，时间复杂度$O(n\log n)$

$code$

#include<cstdio>
#include<cstring>
#include<algorithm>
#define ll long long
#define siz(x) (ed[x]-dfn[x]+1)
using namespace std;
const ll N=3e5+10,M=6e6+10;
struct node{
    ll to,next;
}a[N*2];
ll n,q,tot,cnt,D,ls[N],dep[N];
ll dfn[N],ed[N],rt[N],rfn[N];
void addl(ll x,ll y){
    a[++tot].to=y;a[tot].next=ls[x];ls[x]=tot;
}
void dfs(ll x,ll fa){
    dfn[x]=++cnt;rfn[cnt]=x;dep[x]=dep[fa]+1;D=max(D,dep[x]);for(ll i=ls[x];i;i=a[i].next){
    ll y=a[i].to;if(y==fa)continue;dfs(y,x);}ed[x]=cnt;
}
struct Seg_Tree{
    ll cnt,sum[M],ls[M],rs[M];ll Change(ll x,ll L,ll R,ll pos,ll val){
    ll y=++cnt;sum[y]=sum[x]+val;if(L==R){
    return y;}ll mid=(L+R)>>1;if(pos<=mid)ls[y]=Change(ls[x],L,mid,pos,val),rs[y]=rs[x];else rs[y]=Change(rs[x],mid+1,R,pos,val),ls[y]=ls[x];return y;}ll Ask(ll x,ll y,ll L,ll R,ll l,ll r){
    if(!(sum[y]-sum[x]))return 0;if(L==l&&R==r)return sum[y]-sum[x];ll mid=(L+R)>>1;if(r<=mid)return Ask(ls[x],ls[y],L,mid,l,r);if(l>mid)return Ask(rs[x],rs[y],mid+1,R,l,r);return Ask(ls[x],ls[y],L,mid,l,mid)+Ask(rs[x],rs[y],mid+1,R,mid+1,r);}
}T;
int main()
{
    scanf("%lld%lld",&n,&q);for(ll i=1;i<n;i++){
    ll x,y;scanf("%lld%lld",&x,&y);addl(x,y);addl(y,x);}dfs(1,0);for(ll i=1;i<=n;i++){
    int x=rfn[i];rt[i]=T.Change(rt[i-1],1,D,dep[x],siz(x)-1);}while(q--){
    ll p,k;scanf("%lld%lld",&p,&k);ll ans=min(k,dep[p]-1)*(siz(p)-1);printf("%lld\n",ans+T.Ask(rt[dfn[p]],rt[ed[p]],1,D,dep[p]+1,dep[p]+k));}
}