POJ 1741 Tree 树的点分治

Tree

poj 1741

#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<vector>
using namespace std;
const int maxn = 10100;
typedef struct node{
    int to,val;
}edge;
vector<edge> way[maxn];
vector<int> deep;
int siz[maxn],mu[maxn],d[maxn];
bool vis[maxn];
int n,k,ans;
///初始化
void init()
{
    for(int i=0;i<maxn;i++) way[i].clear();
    memset(vis,false,sizeof(vis));
    mu[0] = 0x3f3f3f3f;
    ans = 0;
}
///加边
void add(int u,int v,int val)
{
    edge temp = {v,val};
    way[u].push_back(temp);
    temp = {u,val};
    way[v].push_back(temp);
}
///求树重心;
void dfsize(int u,int fa)
{
    siz[u] = 1;
    mu[u] = 0;
    for(int i=0;i<way[u].size();i++)
    {
        int t = way[u][i].to;
        if(t == fa || vis[t] ) continue;
        dfsize(t,u);
        siz[u] += siz[t];
        mu[u] = max(mu[u],siz[t]);
    }
}
void dfsroot(int u,int fa,int allnode,int &root)
{
    mu[u] = max(mu[u],allnode - siz[u]);
    if(mu[u] < mu[root]) root = u;
    for(int i=0;i<way[u].size();i++)
    {
        int t = way[u][i].to;
        if(t!=fa && !vis[t]) dfsroot(t,u,allnode,root);
    }
}
int getroot(int u,int fa)
{
    int root = 0;
    dfsize(u,fa);
    dfsroot(u,fa,siz[u],root);
    return root;
}
///
void getdeep(int x,int fa)
{
    deep.push_back(d[x]);
    for(int i=0;i<way[x].size();i++)
    {
        int t = way[x][i].to;
        if(t == fa || vis[t]) continue;
        d[t] = d[x] + way[x][i].val;
        getdeep(t,x);
    }
}
int cal(int x,int now)
{
    d[x] = now; deep.clear();
    getdeep(x,0);
    sort(deep.begin(),deep.end());
    int all = 0;
    for(int i=0,j=deep.size()-1;i<j;){
        if(deep[i] + deep[j] <= k) {all += j-i;i++;}
        else j--;
    }
    return all;
}
void solve(int x)
{
    vis[x] = true;
    ans += cal(x,0);
    for(int i=0;i<way[x].size();i++)
    {
        int t = way[x][i].to;
        if(vis[t]) continue;
        ans -= cal(t,way[x][i].val);
        solve(getroot(t,x));
    }
}
int main()
{
    int u,v,tag;
    while(~scanf("%d%d",&n,&k),n+k)
    {
        init();
        for(int i=1;i<n;i++){
            scanf("%d%d%d",&u,&v,&tag);
            add(u,v,tag);
        }
        int root = getroot(1,0);
        solve(root);
        printf("%d\n",ans);
    }
}