1. 程式人生 > >POJ 1741 Tree 樹上點分治

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 '\n'
#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\n";
	fclose(stdin);fclose(stdout);system("out.txt");
#endif
	return 0;
}