题解：CF2063E Triangle Tree

_wsq_ · 2025-01-23 09:01:23 · 题解

求所有点对的 f 之和，一眼树上启发式合并。

三角形两边之和大于第三边，两边之差小于第三边。设 a=\min(\operatorname{dist}(u,\operatorname{lca}( u,v)),\operatorname{dist}(v,\operatorname{lca}( u,v))),b=\max(\operatorname{dist}(u,\operatorname{lca}( u,v)),\operatorname{dist}(v,\operatorname{lca}( u,v)))，则答案显然是 (a+b)-(a-b)-1=a+b-a+b-1=2b-1。因此我们可以用平衡树维护每个节点到当前节点的距离和平衡树内值在某一段的数之和。那么增加答案时就直接用新加进来的值把平衡树分裂为两部分，一部分用这部分的和更新，另一部分用新加进来的值和这部分的大小更新（这个式子比较复杂，细节见代码第 71 行）。然后还要注意一定要每个子树都更新完答案再统一加入平衡树（否则 \operatorname{lca} 就不是当前节点了），以及重儿子回溯时要将平衡树里所有的节点都加 1（可以单开一个变量表示总体的修改）。

代码：

#include <iostream>
#include <cstdlib>
using namespace std;
#define maxn 300005
#define int long long
struct edge{
    int to,next;
}e[maxn<<1];
struct node{
    int l,r,siz,pri,val,sum;
}fhq[maxn];
int n,h[maxn],tot,totf,root,sz[maxn],hson[maxn],ans,cnt;
void addedge(int u,int v){
    e[++tot].to=v;
    e[tot].next=h[u];
    h[u]=tot;
    return;
}
void pushup(int x){
    fhq[x].siz=fhq[fhq[x].l].siz+fhq[fhq[x].r].siz+1;
    fhq[x].sum=fhq[fhq[x].l].sum+fhq[fhq[x].r].sum+fhq[x].val;
    return;
}
int newnode(int x){
    fhq[++totf].siz=1;
    fhq[totf].val=x;
    fhq[totf].sum=x;
    fhq[totf].pri=rand();
    return totf;
}
void split(int num,int rt,int &x,int &y){
    if(!rt){
        x=y=0;
        return;
    }
    if(fhq[rt].val<=num){
        x=rt;
        split(num,fhq[x].r,fhq[x].r,y);
    }
    else{
        y=rt;
        split(num,fhq[y].l,x,fhq[y].l);
    }
    pushup(rt);
    return;
}
int merge(int x,int y){
    if(!x||!y){
        return x+y;
    }
    if(fhq[x].pri<fhq[y].pri){
        fhq[x].r=merge(fhq[x].r,y);
        pushup(x);
        return x;
    }
    else{
        fhq[y].l=merge(x,fhq[y].l);
        pushup(y);
        return y;
    }
}
void add(int num){
    int x,y;
    split(num-cnt,root,x,y);
    root=merge(x,merge(newnode(num-cnt),y));
    return;
}
int query(int num){
    int x,y,ret;
    split(num-cnt,root,x,y);
    ret=((fhq[x].sum+cnt*fhq[x].siz)<<1)-fhq[x].siz+fhq[y].siz*((num<<1)-1);//这是代码第 71 行
    root=merge(x,y);
    return ret;
}
void clear(int x){
    if(!x){
        return;
    }
    clear(fhq[x].l);
    clear(fhq[x].r);
    fhq[x].l=0;
    fhq[x].r=0;
    fhq[x].siz=0;
    fhq[x].val=0;
    fhq[x].pri=0;
    fhq[x].sum=0;
    return;
}
void dfs(int x,int fa){
    sz[x]=1;
    hson[x]=0;
    for(int i=h[x];i;i=e[i].next){
        if(e[i].to==fa){
            continue;
        }
        dfs(e[i].to,x);
        sz[x]+=sz[e[i].to];
        if(sz[e[i].to]>sz[hson[x]]){
            hson[x]=e[i].to;
        }
    }
}
void queryl(int x,int fa,int dis){
    ans+=query(dis);
    for(int i=h[x];i;i=e[i].next){
        if(e[i].to==fa){
            continue;
        }
        queryl(e[i].to,x,dis+1);
    }
    return;
}
void addl(int x,int fa,int dis){
    add(dis);
    for(int i=h[x];i;i=e[i].next){
        if(e[i].to==fa){
            continue;
        }
        addl(e[i].to,x,dis+1);
    }
    return;
}
void dsu(int x,int fa,bool typ){
    for(int i=h[x];i;i=e[i].next){
        if(e[i].to==fa||e[i].to==hson[x]){
            continue;
        }
        dsu(e[i].to,x,false);
    }
    if(hson[x]){
        dsu(hson[x],x,true);
    }
    for(int i=h[x];i;i=e[i].next){
        if(e[i].to==fa||e[i].to==hson[x]){
            continue;
        }
        queryl(e[i].to,x,1);
        addl(e[i].to,x,1);
    }
    add(0);
    cnt++;
    if(!typ){
        clear(root);
        root=0;
        totf=0;
        cnt=0;
    }
    return;
}
signed main(){
    int t;
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cin>>t;
    while(t--){
        cin>>n;
        for(int i=1;i<n;i++){
            int u,v;
            cin>>u>>v;
            addedge(u,v);
            addedge(v,u);
        }
        dfs(1,0);
        dsu(1,0,false);
        cout<<ans<<'\n';
        for(int i=1;i<=n;i++){
            h[i]=0;
            sz[i]=0;
        }
        tot=0;
        ans=0;
    }
    return 0;
}