bzoj 1468: Tree 点分治
十分巧妙的一种数据结构
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;
}