cqyz oj | 树的分治 | 树形DP | 树的重心

Description

　　给定一棵N个节点的带权树，定义dist(u,v)为u,v两点间的最短路径长度，路径的长度义为路径上所有边的权和。再给定一个K，如果对于不同的两个结点a,b，如果满足dist(a,b)<=K，则称(a,b)为合法点对。
　　你的任务是求合法点对个数。

Input

　　第一行包含两个个整数N和K，接下来的N-1行，每行包含三个整数：u,v,len，表示树边(u,v)的长度len。

Output

　　一个整数，表示合法点对的数目。

Sample Input 1

5 4
1 2 3
1 3 1
1 4 2
3 5 1

Sample Output 1

8

Hint

0<N<100 000,0<K<1000000000,0<len<10000

---

看到数据范围先想到用long long
直接跳过O(n^2)暴力算法

满足条件的点a和b有以下两种情况
一、a,b之间的路径经过根，即a,b在根节点的两棵不同子树里，或者a,b其中之一是根节点
1.计算root到各点距离dis[j]
2.在dis[]中统计满足dis[a]+dis[b]<=K的点对数量X，转化为化装晚会一题
3.去掉来自同一子树中的点对数Y
答案为X-Y
二、a,b在之间的路径不经过root，即a,b在同一棵子树里
这种情况就是分治算法的子问题，可以递归求
（为什么要先X-Y再在子树里重新求呢？因为由“一、”里计算出来的点对本来可以不经过root，相当于多算了一段长度，导致有些原本满足条件的点对因为这截多算的长度而未被统计到）

由这个思路，设如下函数：
getdist(i)计算以i为根的子树中i到各点的距离 (O(size[i]))
F(i,k)计算以i为根的子树中dis[a]+dis[b]<=k的点对数量(O(size[i]log_2size[i]))
calc(i)递归函数 (O(n))
该方法时间复杂度(O(depth * nlogn))，如果随便找一个根可能导致depth为n，时间退化为(O(n^2logn))

要使depth最小，想到找树的重心，因此又设一个函数getroot(i)找i所在的子树里的重心为根，(O(size[i]))
这样时间稳定在(O(n*(log_2n)^2))

此题代码实现对于我这种蒟蒻有一定难度，有很多函数调用之间的细节：

#include<cstdio>
#include<iostream>
#include<queue>
#include<cstring>
#include<algorithm>
#include<cmath>
#define maxn 100005 
using namespace std;
typedef long long ll;
int fir[maxn],ne[maxn*2],to[maxn*2],w[maxn*2],newp=0;
void add(int x,int y,int z){
	ne[++newp]=fir[x];
	fir[x]=newp;
	to[newp]=y;
	w[newp]=z;
}
int siz[maxn],mn,rt;
bool mark[maxn]={0};
void getroot(int u,int f,int tot){
	siz[u]=1;
	int tmp=0;
	for(int i=fir[u];i;i=ne[i]){
		int v=to[i];
		if(v==f || mark[v])continue;
		getroot(v,u,tot);
		siz[u]+=siz[v];
		tmp=max(tmp,siz[v]);
	}
	tmp=max(tmp,tot-siz[u]);
	if(tmp<mn)rt=u,mn=tmp;
}

ll dis[maxn];
int cnt;
void getdist(int u,int f,ll dist){
	dis[cnt++]=dist;
	for(int i=fir[u];i;i=ne[i]){
		int v=to[i];
		if(v==f || mark[v])continue;
		getdist(v,u,dist+w[i]);
	}
}
ll F(int u,int k){
	cnt=0;
	getdist(u,0,0);
	sort(dis,dis+cnt);
	ll res=0;
	for(int i=0;i<cnt;i++)
		res+=upper_bound(dis,dis+i,k-dis[i])-dis;
	return res;
}
int n,K;
ll calc(int u){
	ll t=F(u,K);
	mark[u]=1;
	for(int i=fir[u];i;i=ne[i]){
		int v=to[i];
		if(mark[v])continue;
		t-=F(v,K-2*w[i]);
		mn=n;
		getroot(v,0,siz[v]);
		t+=calc(rt);
	}
	return t;
}
int main(){
	scanf("%d%d",&n,&K);mn=n;
	for(int i=1,x,y,z;i<n;i++){
		scanf("%d%d%d",&x,&y,&z);
		add(x,y,z);add(y,x,z);
	}
	getroot(1,0,n);
	printf("%lld",calc(rt));
	return 0;
}