P4178 Tree

题目

P4178 Tree

统计树上距离小于等于 \(k\) 的点对数量。

分析

这里使用容斥法进行计算，当然也可以使用树状数组来解决。

首先我们可以求出当前点所有向下路径两两拼接得到的小于等于 \(k\) 的路径个数，这个可以排序过后双指针扫一遍即可。

但是我们发现这样会有不合法的情况出现，也就是两条路径来自同一个子树，于是可以考虑容斥掉这样的路径。

照这样看，我们直接分别减掉每一个子树单独计算的答案就行了，但是其实我们多算的不是这些路径，多算的路径其实是：原权值 \(+val_{(u,v)}\) 的所有路径。

那么我们可以在 \(Calc\) 传参的时候考虑直接初始化距离为 \(val_{(u,v)}\) 即可，相当于给所有路径都加了一个这个。

然后对每一层分治中心都这样做即可。

代码

#include<bits/stdc++.h>
using namespace std;
template <typename T>
inline void read(T &x){
    x=0;char ch=getchar();bool f=false;
    while(!isdigit(ch)) f|=ch=='-',ch=getchar();
    while(isdigit(ch)) x=(x<<1)+(x<<3)+(ch^48),ch=getchar();
    x=f?-x:x;return;
}
template <typename T>
inline void write(T x){
    if(x<0) putchar('-'),x=-x;
    if(x>9) write(x/10);
    putchar(x%10^48);return;
}
#define ll long long
const int N=1e5+5,M=1e7+5;
int n,m,Ans;
int nex[N],head[N],to[N],val[N],idx;
bool vis[N],Vis[N];
int siz[N],path[N];
int path_cnt,MaxRoot,Root,Size;
inline void add(int u,int v,int w){
	nex[++idx]=head[u];
	to[idx]=v;
	head[u]=idx;
	val[idx]=w;
	return ; 
} 
void FindRoot(int x){
	vis[x]=true;siz[x]=1;
	int Max=0;
	for(int i=head[x];i;i=nex[i]){
		int y=to[i];
		if(vis[y]||Vis[y]) continue;
		FindRoot(y);siz[x]+=siz[y];
		Max=max(Max,siz[y]);
	} 
	Max=max(Max,Size-siz[x]);
	if(Max<MaxRoot) MaxRoot=Max,Root=x;
	vis[x]=false;
	return ;
}
void GetPath(int x,int D){
	path[++path_cnt]=D;vis[x]=true;
	for(int i=head[x];i;i=nex[i]){
		int y=to[i];
		if(vis[y]||Vis[y]) continue;
		GetPath(y,D+val[i]);
	}
	vis[x]=false;
	return ;
}

int Calc(int x,int D){
	path_cnt=0;
	GetPath(x,D);
	int r=path_cnt,sum=0;
	sort(path+1,path+path_cnt+1);
	for(int l=1;l<=path_cnt;l++){
		while(r>=1&&path[l]+path[r]>m) r--;
		if(r<l) break;
		sum+=r-l+1;
	}
	return sum;
}
void DFS(int x){
	Ans+=Calc(x,0);Vis[x]=true;
	for(int i=head[x];i;i=nex[i]){
		int y=to[i];
		if(Vis[y]) continue;
		Ans-=Calc(y,val[i]);
		Root=0,MaxRoot=n,Size=siz[y],FindRoot(y),DFS(Root);
	}
	Vis[x]=false;
	return ;
}
signed main(){
	read(n);
	for(int i=1;i<n;i++){
		int u,v,w;
		read(u),read(v),read(w);
		add(u,v,w);
		add(v,u,w);
	}	
	read(m);
	Root=0,MaxRoot=n,Size=n,FindRoot(1),DFS(Root);
	write(Ans-n);
	return 0;
}