bzoj 1468: Tree 点分治

Posted 2021-12-02 guangheli

十分巧妙的一种数据结构

Code:

#include<bits/stdc++.h>
#define N 400010 
#define setIO(s) freopen(s".in","r",stdin) 
using namespace std;
int to[N<<1],head[N<<1],nex[N<<1],val[N<<1];
int f[N],root,m,dep[N],siz[N],sn,d[N],tot,ans,cnt; 
int vis[N];
void add(int u,int v,int c){ nex[++cnt]=head[u],head[u]=cnt,to[cnt]=v,val[cnt]=c;  }
//求重心
void get(int x,int fa){
    f[x]=0,siz[x]=1; 
    for(int v=head[x];v;v=nex[v]){
        if(to[v]==fa||vis[to[v]]) continue; 
        get(to[v],x); 
        siz[x]+=siz[to[v]]; 
        f[x]=max(f[x],siz[to[v]]); 
    }      
    f[x]=max(f[x],sn-siz[x]); 
    if(f[x]<f[root]) root=x; 
}
void getdis(int x,int fa){
    d[++tot]=dep[x];
    for(int v=head[x];v;v=nex[v]){
        if(to[v]==fa||vis[to[v]]) continue; 
        dep[to[v]]=dep[x]+val[v],getdis(to[v],x); 
    }
}
int calc(int x){     
    tot=0; 
    getdis(x,0); 
    sort(d+1,d+tot+1);  
    int i=1,j=tot,sum=0; 
    while(i<j) {
        if(d[i]+d[j]<=m) {
            sum+=j-i,i++;  
        }
        else j--;              
    }
    return sum; 
}
void dfs(int x){
    dep[x]=0; 
    vis[x]=1; 
    ans+=calc(x); 
    for(int v=head[x];v;v=nex[v]){
        if(!vis[to[v]]) {
            dep[to[v]]=val[v]; 
            ans-=calc(to[v]);          
            sn=siz[to[v]]; 
            root=0; 
            get(to[v],0); 
            dfs(root); 
        }
    }
}
int main(){
    // setIO("input"); 
    int n; 
    scanf("%d",&n); 
    for(int i=1;i<n;++i) 
    {
        int x,y,z; 
        scanf("%d%d%d",&x,&y,&z); 
        add(x,y,z),add(y,x,z); 
    }
    cnt=0,ans=0; 
    scanf("%d",&m); 
    f[0]=0x7f7f7f7f; 
    sn=n; 
    root=0,get(1,0),dfs(root);
    printf("%d",ans); 
    return 0; 
}