poj 1741 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

题意:给你一棵树，问你树上路径长度小于等于k的路径数量

思路:点分治模板题

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<iostream>
#include<string>
#include<vector>
#include<stack>
#include<bitset>
#include<cstdlib>
#include<cmath>
#include<set>
#include<list>
#include<deque>
#include<map>
#include<queue>
#define ll long long int
using namespace std;
inline ll gcd(ll a,ll b){return b?gcd(b,a%b):a;}
inline ll lcm(ll a,ll b){return a/gcd(a,b)*b;}
int moth[13]={0,31,28,31,30,31,30,31,31,30,31,30,31};
int dir[4][2]={1,0 ,0,1 ,-1,0 ,0,-1};
int dirs[8][2]={1,0 ,0,1 ,-1,0 ,0,-1, -1,-1 ,-1,1 ,1,-1 ,1,1};
const int inf=0x3f3f3f3f;
const ll mod=1e9+7;
int head[10007],vis[10007],d[10007];
struct node{
    int to,v,next;
};
node edge[10007<<1];
int cnt,ans,n,k;
void init(){
    cnt=0;
    ans=0;
    memset(head,0,sizeof(head));
    memset(vis,0,sizeof(vis));
}
void add(int from,int to,int val){
    edge[++cnt].v=val;
    edge[cnt].to=to;
    edge[cnt].next=head[from];
    head[from]=cnt;
}
int son[10007]; //儿子节点数 
int now_size;  //子树的最小节点数，用于找重心 
int sz;   //当前操作树的节点数 
int root; //当前根节点 
void find_root(int u,int fa){ //找树重心 
    son[u]=1; int res=0;
    for(int i=head[u];i;i=edge[i].next){
        if(vis[edge[i].to]||edge[i].to==fa) continue;
        int to=edge[i].to;
        find_root(to,u);
        son[u]+=son[to];
        res=max(res,son[to]);
    }
    res=max(res,sz-son[u]);
    if(res<now_size) now_size=res,root=u;
}
int a[10007]; //临时储存路径值 
int tot;
void get_dis(int u,int fa){ //计算d数组 d[i]表示i点距离当前根节点的距离 
    a[++tot]=d[u];
    for(int i=head[u];i;i=edge[i].next){
        if(vis[edge[i].to]||edge[i].to==fa) continue;
        int to=edge[i].to;
        d[to]=d[u]+edge[i].v;
        get_dis(to,u);
    }
}
int solve(int u,int dis){ //找符合情况的个数 
    d[u]=dis; tot=0;
    get_dis(u,u);
    sort(a+1,a+1+tot);
    int l=1; int r=tot; int res=0;
    for(;l<r;++l){
        while(l<r&&a[l]+a[r]>k) --r;
        if(l<r) res+=(r-l);
    }
    return res;
}
void dfs(int u){ //分治 
    vis[u]=1;
    ans+=solve(u,0);
    int totsz=sz;
    for(int i=head[u];i;i=edge[i].next){
        if(vis[edge[i].to]) continue;
        int to=edge[i].to;
        ans-=solve(to,edge[i].v); //容斥思想 
        now_size=inf; root=0; 
        sz=son[to]>son[u]?totsz-son[u]:son[to];
        find_root(to,0);
        dfs(root);
    }
}
int main(){
    ios::sync_with_stdio(false);
    while(cin>>n>>k){
        if(!n&&!k) break;
        init();
        for(int i=1;i<n;i++){
            int from,to,val;
            cin>>from>>to>>val;
            add(from,to,val); add(to,from,val);
        }
        now_size=inf,sz=n,root=0;
        find_root(1,0); //找重心 
        dfs(root);
        cout<<ans<<endl;
    }
    return 0;
}