【树分治】poj1741 tree

题目

Description

Give a tree with n vertices,each edge has a length(positive integer less than 1001).

Define dist(u,v)=The min distance between node u and v.

Give an integer k,for every pair (u,v) of vertices is called valid if and only if dist(u,v) not exceed k.

Write a program that will count how many pairs which are valid for a given tree.

Input

The input contains several test cases. The first line of each test case contains two integers n, k. (n<=10000) The following n-1 lines each contains three integers u,v,l, which means there is an edge between node u and v of length l.

The last test case is followed by two zeros.

Output

For each test case output the answer on a single line.

Sample Input

5 4

1 2 3

1 3 1

1 4 2

3 5 1

0 0

Sample Output

8

题目大意

给定无根树一颗，求满足u到v的距离小于等于k的（u，v）点对个数。

男人八题之一，点分治。

点分治，在一颗无根树中找到树的重心（重心就是树中的一点，满足这一点的所有子节点中size值最大的最小），然后把重心去掉，对于重心的每一个子树进行相同操作。

这道题就是对于每一个子树，统计一下在根节点的不同子树中的两点对个数即可（相同子树中的点对分治到子树中操作即可）

先建图：

struct edge{
int v,w,next;
}e[MAXM];
void add_edge(int u, int v, int w)
{
e[edge].v = v;
e[edge].w = w;
e[edge].next = head[u];
head[u] = edge++;
}

dfs找size：

void dfssize(int u,int fa){
size[u]=1;
mx[u]=0;
for(int i=head[u];i!=-1;i=e[i].next){
int v=e[i].v;
if(v!=fa&&!vis[v]){
dfssize(v,u);
size[u]+=size[v];
if(size[v]>mx[u])
mx[u]=size[v];
}
}
}

然后dfs找root：

void dfsroot(int r,int u,int fa){
if(size[r]-size[u]>mx[u])
mx[u]=size[r]-size[u];
if(mx[u]<minn)
minn=mx[u],root=u;
for(int i=head[u];i!=-1;i=e[i].next){
int v=e[i].v;
if(v!=fa&&!vis[v])
dfsroot(r,v,u);
}
}

然后再写一个dfs求两点距离：

void dfsdis(int u,int d,int fa){
dis[num++]=d;
for(int i=head[u];i!=-1;i=e[i].next){
int v=e[i].v;
if(v!=fa&&!vis[v])
dfsdis(v,d+e[i].w,u);
}
}

然后写主dfs分治：

void dfsroot(int r,int u,int fa){
if(size[r]-size[u]>mx[u])
mx[u]=size[r]-size[u];
if(mx[u]<minn)
minn=mx[u],root=u;
for(int i=head[u];i!=-1;i=e[i].next){
int v=e[i].v;
if(v!=fa&&!vis[v])
dfsroot(r,v,u);
}
}

还有一个是计算刚才说的跨过根节点的点对个数：

int calc(int u,int d){
int res=0;
num=0;
dfsdis(u,d,0);
sort(dis,dis+num);
int i=0,j=num-1;
while(i<j){
while(dis[i]+dis[j]>k&&i<j) j--;
res+=j-i;
i++;
}
return res;
}

然后就完事了。。。对于我这个dfs写的很low的人，写四个dfs也是醉了。

直接贴全代码：

#include <iostream>
#include <cstdio>
#include <vector>
#include <climits>
#include <cstring>
#include <algorithm>
#define MAXN 20000
#define MAXM 50000
using namespace std;
struct edge{
int v,w,next;
}e[MAXM];
int edge,ans,num,minn;
int size[MAXN],dis[MAXN],mx[MAXN],head[MAXN];
bool vis[MAXN];
int n,k;
int root;
void pre(){
memset(head,-1,sizeof head);
memset(vis,0,sizeof vis);
edge=ans=0;
}
void add_edge(int u, int v, int w)
{
e[edge].v = v;
e[edge].w = w;
e[edge].next = head[u];
head[u] = edge++;
}
void dfssize(int u,int fa){
size[u]=1;
mx[u]=0;
for(int i=head[u];i!=-1;i=e[i].next){
int v=e[i].v;
if(v!=fa&&!vis[v]){
dfssize(v,u);
size[u]+=size[v];
if(size[v]>mx[u])
mx[u]=size[v];
}
}
}
void dfsroot(int r,int u,int fa){
if(size[r]-size[u]>mx[u])
mx[u]=size[r]-size[u];
if(mx[u]<minn)
minn=mx[u],root=u;
for(int i=head[u];i!=-1;i=e[i].next){
int v=e[i].v;
if(v!=fa&&!vis[v])
dfsroot(r,v,u);
}
}
void dfsdis(int u,int d,int fa){
dis[num++]=d;
for(int i=head[u];i!=-1;i=e[i].next){
int v=e[i].v;
if(v!=fa&&!vis[v])
dfsdis(v,d+e[i].w,u);
}
}
int calc(int u,int d){
int res=0;
num=0;
dfsdis(u,d,0);
sort(dis,dis+num);
int i=0,j=num-1;
while(i<j){
while(dis[i]+dis[j]>k&&i<j) j--;
res+=j-i;
i++;
}
return res;
}
void dfs(int u){
minn=n;
dfssize(u,0);
dfsroot(u,u,0);
ans+=calc(root,0);
vis[root]=1;
for(int i=head[root];i!=-1;i=e[i].next){
int v=e[i].v;
if(!vis[v]){
ans-=calc(v,e[i].w);
dfs(v);
}
}
}
int main()
{
while(scanf("%d%d",&n,&k)!=EOF)
{
if(!n && !k) break;
pre();
int u, v, w;
for(int i=0;i<n-1;i++)
{
scanf("%d%d%d", &u, &v, &w);
add_edge(u, v, w);
add_edge(v, u, w);
}
dfs(1);
printf("%d\n", ans);
}
return 0;
}  

原po 传送门