【树分治】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 传送门