poj 1741 Tree(树的点分治)

时间:2022-10-04 04:24:02

poj 1741 Tree(树的点分治)

给出一个n个结点的树和一个整数k,问有多少个距离不超过k的点对。

首先对于一个树中的点对,要么经过根结点,要么不经过。所以我们可以把经过根节点的符合点对统计出来。接着对于每一个子树再次运算。如果不用点分治的技巧,时间复杂度可能退化成\(O(n^2)\)(链)。如果对于子树重新选根,找到树的重心,就一定可以保证时间复杂度在\(O(nlogn)\)内。

具体技巧是:首先选出树的重心,将重心视为根。接着计算出每个结点的深度,以此统计答案。由于子树中可能出现重复情况,需要在子树中相应的减去一部分ans。具体实现在代码中。

#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;

const int maxn=1e4+5;

struct Graph{
struct Edge{
int to, next, v; Graph *bel;
Edge& operator ++(){
return *this=bel->edge[next]; }
}edge[maxn*2];
int cntedge, fir[maxn];
void addedge(int x, int y, int v){
Edge &e=edge[++cntedge];
e.to=y; e.next=fir[x]; e.v=v;
fir[x]=cntedge; e.bel=this;
}
Edge& getlink(int x){ return edge[fir[x]]; }
void RESET(){ cntedge=0; memset(fir, 0, sizeof(fir)); }
}g;

int n, k, size[maxn], w[maxn], dep[maxn];
int cnt[maxn], tail;
bool done[maxn];

//获取子树大小
int getsize(int now, int par, int num){
size[now]=0; w[now]=0;
Graph::Edge e=g.getlink(now);
for (; e.to; ++e){
if (e.to==par||done[e.to]) continue;
size[now]+=getsize(e.to, now, num);
w[now]=max(w[now], size[e.to]);
}
++size[now]; w[now]=max(w[now], num-size[now]);
return size[now];
}

//获取根的位置
int getroot(int now, int par){
Graph::Edge e=g.getlink(now);
int root=now, tmp=now;
for (; e.to; ++e){
if (e.to==par||done[e.to]) continue;
tmp=getroot(e.to, now);
if (w[tmp]<w[root]) root=tmp; //mdzzle
}
return root;
}

//获取子树中结点深度,并创造深度数组
void getdep(int now, int par, int step){
Graph::Edge e=g.getlink(now);
for (; e.to; ++e)
if (e.to!=par&&!done[e.to])
getdep(e.to, now, step+e.v);
dep[now]=step;
cnt[tail++]=step;
}

//通过深度数组统计和小于等于k的数对
int getans(int k, int tail){
sort(cnt, cnt+tail);
--tail; int ans=0;
for (int l=0; l<tail; ++l){
while (cnt[l]+cnt[tail]>k&&l<tail) --tail;
ans+=tail-l;
}
return ans;
}

int solve(int now, int num){
getsize(now, 0, num); //获取当前树的子树大小
if (size[now]==1) return 0;
now=getroot(now, 0); //凭借子树大小找到重心
tail=0; //把深度数组复原
getdep(now, 0, 0); //获取每个结点的深度,构建深度数组
//获取答案,别忘记减掉重复的点对
int ans=getans(k, tail);
done[now]=true; //堵住当前点
Graph::Edge e=g.getlink(now);
for (; e.to; ++e) if (!done[e.to]){
tail=0; getdep(e.to, 0, 0);
ans-=getans(k-e.v*2, tail);
ans+=solve(e.to, size[e.to]);
}
return ans;
}

int main(){
while (~scanf("%d%d", &n, &k)&&n&&k){
int t1, t2, t3; g.RESET();
for (int i=1; i<=n; ++i) done[i]=false;
for (int i=1; i<n; ++i){
scanf("%d%d%d", &t1, &t2, &t3);
g.addedge(t1, t2, t3);
g.addedge(t2, t1, t3);
}
printf("%d\n", solve(1, n));
}
return 0;
}