poj1741(树分治)

链接：点击打开链接

题意：求出一棵树上小于等于K的路径的条数

代码：

#include <map>
#include <set>
#include <queue>
#include <string>
#include <math.h>
#include <vector>
#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <string.h>
#include <algorithm>
using namespace std;
const int siz=10005;
const int INF=0x3f3f3f3f;
int read(){
	int x=0,f=1;char ch=getchar();
	while(ch<'0'||ch>'9'){if(ch=='-')f=-1;ch=getchar();}
	while(ch>='0'&&ch<='9'){x=x*10+ch-'0';ch=getchar();}
	return x*f;
}
struct node{
    int u,v;
};
vector<node> G[siz];
int K,op,rt,ans,sum;
int d[siz],f[siz],dep[siz],son[siz],vis[siz];
void getrt(int s,int fa){
    int i,u;
    son[s]=1,f[s]=0;
    for(i=0;i<G[s].size();i++){
        u=G[s][i].u;
        if(u==fa||vis[u]==1)
        continue;
        getrt(u,s);
        son[s]+=son[u];
        f[s]=max(f[s],son[u]);
    }
    f[s]=max(f[s],sum-son[s]);
    if(f[s]<f[rt])
    rt=s;
}
void getdep(int s,int fa){
    int i,u,v;
    dep[op++]=d[s];
    for(i=0;i<G[s].size();i++){
        u=G[s][i].u;
        v=G[s][i].v;
        if(u==fa||vis[u]==1)
        continue;
        d[u]=d[s]+v;
        getdep(u,s);
    }
}
int cal(int s,int tmp){
    int l,r,res;
    op=res=0,d[s]=tmp;
    getdep(s,-1);
    sort(dep,dep+op);
    l=0,r=op-1;
    while(l<r){
        if(dep[l]+dep[r]<=K){
            res+=(r-l);
            l++;
        }
        else
        r--;
    }
    return res;
}
void dfs(int s){
    int i,u,v;
    ans+=cal(s,0);
    vis[s]=1;
    for(i=0;i<G[s].size();i++){
        u=G[s][i].u;
        v=G[s][i].v;
        if(vis[u])
        continue;
        ans-=cal(u,v);
        rt=0,sum=son[u];
        getrt(u,-1);
        dfs(rt);
    }
}
int main(){                                     //点分治模板
    int i,u,v,w,n;
    while(1){
        n=read(),K=read();
        if(n==0&&K==0)
        break;
        for(i=1;i<=n;i++)
        G[i].clear();
        memset(vis,0,sizeof(vis));
        for(i=1;i<n;i++){
            u=read(),v=read(),w=read();
            G[u].push_back((node){v,w});
            G[v].push_back((node){u,w});
        }
        ans=0,sum=n,f[0]=INF;
        getrt(1,-1);
        dfs(rt);
        printf("%d\n",ans);
    }
    return 0;
}