POJ 1741 Tree 树上点分治

题目链接:http://poj.org/problem?id=1741

题意:

给定一棵包含$n$个点的带边权树,求距离小于等于K的点对数量

题解:

显然,枚举所有点的子树可以获得答案,但是朴素发$O(n^2logn)$算法会超时,

利用树的重心进行点分治可以将$O(n^2logn)$的上界优化为近似$O(nlogn)$

足以在1000ms的测试时间内通过

具体原理参考注释

#include<iostream>
#include<map>
#include<string>
#include<cstring>
#include<vector>
#include<algorithm>
#include<set>
#include<sstream>
#include<cstdio>
#include<cmath>
#include<climits>
#define endl '
'
#define ll long long
#define ull unsigned long long
#define fi first
#define se second
#define mp make_pair
#define pii pair<int,int>
#define all(x) x.begin(),x.end()
#define IO ios::sync_with_stdio(false)
#define rep(ii,a,b) for(int ii=a;ii<=b;++ii)
#define per(ii,a,b) for(int ii=b;ii>=a;--ii)
#define forn(x,i) for(int i=head[x];i;i=e[i].next)
#define show(x) cout<<#x<<"="<<x<<endl
#define showa(a,b) cout<<#a<<'['<<b<<"]="baidu<a[b]<<endl
#define show2(x,y) cout<<#x<<"="<<x<<" "<<#y<<"="<<y<<endl
#define show3(x,y,z) cout<<#x<<"="<<x<<" "<<#y<<"="<<y<<" "<<#z<<"="<<z<<endl
#define show4(w,x,y,z) cout<<#w<<"="<<w<<" "<<#x<<"="<<x<<" "<<#y<<"="<<y<<" "<<#z<<"="<<z<<endl
using namespace std;
const int maxn=1e6+10,maxm=2e6+10;
const int INF=0x3f3f3f3f;
const int mod=1e9+7;
const double PI=acos(-1.0);
//head
int casn,n,m,k;
ll val[maxn],dis[maxn],ans,maxt,dfn;
int deep[maxn],vis[maxn],size[maxn];
int dp[maxn],allnode;
struct node {int to,next;ll cost;}e[maxm];int head[maxn],nume;//静态链表存图
void add(int a,int b,ll c){e[++nume]=(node){b,head[a],c};head[a]=nume;}
int mid;
void init(){//初始化
	memset(head,0,sizeof head);
	memset(dis,0,sizeof dis);
	memset(vis,0,sizeof vis);
	nume=0;
}
void getmid(int now,int pre){//dfs求树的重心
	size[now]=1;//当前点为根,其子树的节点数
	for(int i=head[now];i;i=e[i].next){
		int to=e[i].to;
		if(to==pre||vis[to]) continue;
		getmid(to,now);//递归计算子树的大小
		size[now]+=size[to];
	}
	dp[now]=max(size[now],allnode-size[now]);//dp[i]表示以i为根建立子树的时候,最大的子树大小
	if(maxt>dp[now]){//maxt为最大的子树大小
		maxt=dp[now];
		mid=now;
	}
}
void dfs(int now,int pre){//计算深度
	deep[++dfn]=dis[now];
	for(int i=head[now];i;i=e[i].next){
		int to=e[i].to,cost=e[i].cost;
		if(to==pre||vis[to]) continue;
		dis[to]=dis[now]+cost;
		dfs(to,now);
	}
}
int cal(int rt,int len){//计算rt为根的子树中,深度之和>=k的点对数量
	dis[rt]=len,dfn=0;
	dfs(rt,0);//以rt为根,dfs计算其子树中所有点的深度
	sort(deep+1,deep+dfn+1);
	int res=0;
	for(int l=1,r=dfn;l<r;){//排序后从两端向中间逼近,总复杂度nlogn
		if(deep[l]+deep[r]<=k){
			res+=r-l;
			l++;
		}else r--;
	}
	return res;
}
void dc(int rt){//分治以rt为根的子树
	vis[rt]=1;
	ans+=cal(rt,0);//初步计算以rt为根子树答案,包含重复情况
	for(int i=head[rt];i;i=e[i].next){
		int to=e[i].to;
		if(vis[to]) continue;
		ans-=cal(to,e[i].cost);//以子节点为根的子树,设置其距离下界为len,对于其子树而言,如果距离减少子树到rt的距离,依然成立的话,必然会被重复计算
		allnode=size[to];
		mid=0,maxt=INF;
		getmid(to,rt);//寻找以rt为根的子树的重心
		dc(mid);//以子树重心为树上点分治的起点,保证总复杂度为n(logn)^2级别
	}

}
int main() {
//#define test
#ifdef test
	auto _start = chrono::high_resolution_clock::now();
	freopen("in.txt","r",stdin);freopen("out.txt","w",stdout);
#endif
	IO;
	while(cin>>n>>k,n+k){
		init();
		int a,b,c;
		rep(i,2,n){
			cin>>a>>b>>c;
			add(a,b,c);
			add(b,a,c);
		}
		mid=ans=0;
		allnode=n,maxt=INF;
		getmid(1,0);
		dc(mid);
		cout<<ans<<endl;
	}
#ifdef test
	auto _end = chrono::high_resolution_clock::now();
  cerr << "elapsed time: " << chrono::duration<double, milli>(_end - _start).count() << " ms
";
	fclose(stdin);fclose(stdout);system("out.txt");
#endif
	return 0;
}