P4178 Tree 题解

题面

使用点分治的思想，对每个子树以中心为根dfs一遍算出 \(b,d\) 数组，然后把 \(a\) 数组（子树内所有点）按照 \(d\) 从小到大排序，用两个指针 \(l,r\) 扫数组。注意到这个题需要算个数，所以还需要考虑把在同一个子树内的删掉。另外开一个数组 \(c_i\) 动态维护 \([l+1,r]\) 区间内在 \(i\) 的子树中的点数，扫一遍即可。复杂度 \(O(n\log^2 n)\)，复杂度与平衡树或者权值线段树的一样，但是好写多了……

点击查看代码

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std;
const int N=4e4+13;
struct Edge{int v,w,nxt;}e[N<<1];
int n,k,tot,h[N],rt,siz[N],maxson[N],a[N],b[N],c[N],d[N],cnt,ans;
bool vis[N];
inline bool cmp(const int &x,const int &y){return d[x]<d[y];}
inline void add(int u,int v,int w){e[++tot]=(Edge){v,w,h[u]};h[u]=tot;}
void getroot(int u,int fa,int total){
	siz[u]=1,maxson[u]=0;
	for(int i=h[u];i;i=e[i].nxt){
		int v=e[i].v;if(v==fa||vis[v]) continue;
		getroot(v,u,total);
		siz[u]+=siz[v];
		maxson[u]=max(maxson[u],siz[v]);
	}
	maxson[u]=max(maxson[u],total-siz[u]);
	if(!rt||maxson[u]<maxson[rt]) rt=u;
}
void init(int u,int fa,int sum,int top){
	a[++cnt]=u,b[u]=top,d[u]=sum;
	for(int i=h[u];i;i=e[i].nxt){
		int v=e[i].v;if(v==fa||vis[v]) continue;
		init(v,u,sum+e[i].w,top);
	}
}
void solve(int u){
	vis[u]=1;a[cnt=1]=u,b[u]=u,d[u]=0;
	for(int i=h[u];i;i=e[i].nxt){
		int v=e[i].v;if(vis[v]) continue;
		init(v,0,e[i].w,v);
	}
	sort(a+1,a+cnt+1,cmp);memset(c,0,sizeof c);
	for(int i=1;i<=cnt;++i) c[b[a[i]]]++;
	for(int l=1,r=cnt;l<=cnt&&l<=r;++l){
		c[b[a[l]]]--;
		while(d[a[l]]+d[a[r]]>k) c[b[a[r]]]--,--r;
		ans+=r-l-c[b[a[l]]];
	}
	for(int i=h[u];i;i=e[i].nxt){
		int v=e[i].v;if(vis[v]) continue;
		rt=0,getroot(v,0,siz[v]);getroot(rt,0,siz[rt]);
		solve(rt);
	}
}
int main(){
	scanf("%d",&n);
	for(int i=1,u,v,w;i<n;++i) scanf("%d%d%d",&u,&v,&w),add(u,v,w),add(v,u,w);
	scanf("%d",&k);
	getroot(1,0,n);
	solve(rt);
	printf("%d\n",ans);
	return 0;
}