点分治 模板 poj1714

点分治 模板 poj1714

点分治

---

Code：

#include <algorithm>
#include <iostream>
#include <cstring>
#include <cstdio>
#include <vector>
#define MP make_pair
#define D(x) cout<<#x<<" = "<<x<<"  "
#define E cout<<endl
using namespace std;
typedef pair<int,int> pii;
const int N = 10005;
const int INF = 0x3f3f3f3f;

int n,K,ans; int sz[N]; bool del[N];
struct Edge{ int to,nxt,w; } e[N<<1];
int head[N], ec=0;
void add(int a,int b,int w){
    ec++; e[ec].to=b; e[ec].w=w; e[ec].nxt=head[a]; head[a]=ec;
}

void getSize(int u,int f){
    sz[u]=1;
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to; if(v==f || del[v]) continue;
        getSize(v,u); sz[u]+=sz[v];
    }
}

pii getRoot(int u,int f,int t){
    pii res=MP(INF,-1); int m=0,s=1;
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to; if(v==f || del[v]) continue;
        res=min(res,getRoot(v,u,t));
        m=max(m,sz[v]); s+=sz[v];
    }
    m=max(m,t-s);
    res=min(res,MP(m,u));
    return res;
}

void getPath(int u,int f,int d,vector<int> &ds){
    ds.push_back(d);
    for(int i=head[u];i;i=e[i].nxt){
        int v=e[i].to; if(v==f || del[v]) continue;
        getPath(v,u,d+e[i].w,ds);
    }
}

int calc(vector<int> &ds){
    int res=0,j=ds.size(); sort(ds.begin(),ds.end());
    for(int i=0;i<ds.size();i++){
        while(j>0 && ds[i]+ds[j-1]>K) --j;
        res+=j-(j>i?1:0);
    }
    return res/2;
}

void solve(int u){
    getSize(u,-1);
    int r=getRoot(u,-1,sz[u]).second;
    del[r]=true;
    for(int i=head[r];i;i=e[i].nxt){
        if(del[e[i].to]) continue;
        solve(e[i].to);
    }
    vector<int> ds; ds.push_back(0);
    for(int i=head[r];i;i=e[i].nxt){
        if(del[e[i].to]) continue;
        vector<int> tp;
        getPath(e[i].to,r,e[i].w,tp);
        ans-=calc(tp);
        ds.insert(ds.end(),tp.begin(),tp.end());
    }
    ans+=calc(ds); del[r]=false;
}

int main(){
    freopen("a.in","r",stdin);
    while(scanf("%d%d",&n,&K)==2){
        if(n==0 || K==0) break;
        memset(head,0,sizeof(head));
        memset(del,false,sizeof(del));
        int a,b,w;
        for(int i=1;i<n;i++){
            scanf("%d%d%d",&a,&b,&w);
            add(a,b,w); add(b,a,w);
        }
        ans=0; solve(1);
        printf("%d\n",ans);
    }
}