[湖南集训]谈笑风生 线段树合并

[湖南集训]谈笑风生 线段树合并

给一棵树\(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;
}