题目链接: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;
}