洛谷P3899：谈笑风生（树上差分+BIT）

题面
题意：一棵树，每次给出一个点p和长度k
问有多少对(b,c)使得p,b是c的祖先，p,c距离小于k

发现只要知道p的子树内，与p距离小于k的点的siz，就可以计算答案
我比较懒，用树上差分+BIT实现

#include <iostream>
#include <fstream>
#include <algorithm>
#include <cmath>
#include <ctime>
#include <cstdio>
#include <cstdlib>
#include <cstring>

using namespace std;
#define mmst(a, b) memset(a, b, sizeof(a))
#define mmcp(a, b) memcpy(a, b, sizeof(b))

typedef long long LL;

const int N=600600;

void read(int &hy)
{
    hy=0;
    char cc=getchar();
    while(cc<'0'||cc>'9')
    cc=getchar();
    while(cc>='0'&&cc<='9')
    {
        hy=(hy<<3)+(hy<<1)+cc-'0';
        cc=getchar();
    }
}

int n,q;
int head[N],lter[N],nex[N<<1],to[N<<1],cnt;
int siz[N],dep[N],fa[N];
int p[N],k[N];
LL bit[N];
LL ans[N];

void add(int u,int v)
{
    to[++cnt]=v;
    nex[cnt]=head[u];
    head[u]=cnt;
} 

void update(int x,LL ad)
{
    x+=3;
    for(;x<=n+10;x+=x&-x)
    bit[x]+=ad;
}

LL sum(LL x)
{
    x+=3;
    LL res=0;
    for(;x;x-=x&-x)
    res+=bit[x];
    return res;
}

void dfs(int x)
{
    siz[x]=1;
    for(int h=head[x];h!=lter[x];h=nex[h])
    ans[to[h]]-=sum(min(n,dep[x]+k[to[h]]));

    for(int h=lter[x];h;h=nex[h])
    if(to[h]!=fa[x])
    {
        dep[to[h]]=dep[x]+1;
        fa[to[h]]=x;
        dfs(to[h]);
        siz[x]+=siz[to[h]];
    }
    update(dep[x],siz[x]-1);

    for(int h=head[x];h!=lter[x];h=nex[h])
    ans[to[h]]+=sum(min(n,dep[x]+k[to[h]]))-siz[x]+1;
}

int main()
{
    cin>>n>>q;
    for(int i=1;i<n;i++)
    {
        int u,v;
        read(u);
        read(v);
        add(u,v);
        add(v,u);
    }
    mmcp(lter,head);
    for(int i=1;i<=q;i++)
    {
        read(p[i]);
        read(k[i]);
        add(p[i],i);
    }
    dfs(1);
    for(int i=1;i<=q;i++)
    {
        ans[i]+=(LL)(siz[p[i]]-1)*(LL)min(dep[p[i]],k[i]);
        printf("%lld\n",ans[i]);
    }

    return 0;
}