[BZOJ 3653] 谈笑风生

题目

给定一棵 n 个点的有根树，另有 q 次询问，每次询问给定
a、k，求有多少个点对 (b, c) 满足 a、b、c 两两不同，a、b 都是
c 的祖先且 a、b 间距离不超过 k。 n, q ≤ 3 × 10^5。

思路

有两种情况。当 b 是 a 的祖先时，b 在 a 的 1 至 k 级祖先
中任选，c 在 a 的子树内任选。
当 b 在 a 子树内时，要求 depb ≤ depa + k，c 在 b 的子树
内任选。可以用主席树（或线段树合并）维护 a 子树内、dep 在
某个区间限制内的 size 和。

代码

#include 
#include 
#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;
}