[POJ1741]Tree

原题地址

点分治经典题，求一棵树上不大于k的路径条数(原题有些小小的不同)

AC code:

#include <cstdio>
#include <vector>
#include <cstring>
#include <algorithm>
using namespace std;
const int N=10010;
const int INF=1<<28;
int  n,k,ans,m,mins,root;
int  son[N];
bool del[N];
vector<int> rd[N];

struct Data{
    int v,w;

    Data(int v,int w):v(v),w(w) {}
};
vector<Data> G[N];

bool read(){
    scanf("%d%d",&n,&k);
    if(n==0&&k==0) return 0;
    for(int i=1;i<n;i++){
        int u,v,w;
        scanf("%d%d%d",&u,&v,&w);
        G[u].push_back(Data(v,w));
        G[v].push_back(Data(u,w));
    }
    return 1;
}

void DFS1(int pre,int p){
    int maxs=-INF;

    son[p]=1;
    for(int i=0;i<G[p].size();i++){
        Data q=G[p][i];
        if(del[q.v]||q.v==pre) continue;
        DFS1(p,q.v);
        son[p]+=son[q.v];
        maxs=max(maxs,son[q.v]);
    }
    maxs=max(maxs,m-son[p]);
    if(maxs<mins){
        mins=maxs;
        root=p;
    }
}

void DFS2(int anc,int pre,int p,int w){
    son[p]=1;
    rd[anc].push_back(w);
    for(int i=0;i<G[p].size();i++){
        Data q=G[p][i];
        if(del[q.v]||q.v==pre) continue;
        DFS2(anc,p,q.v,w+q.w);
        son[p]+=son[q.v];
    }
}

int census(vector<int> a){
    int sum=0,p=0,q=a.size()-1;

    sort(a.begin(),a.end());
    while(1){
        while(q>=0&&a[p]+a[q]>k) q--;
        if(p>q) break;
        sum+=q-p;
        p++;
    }

    return sum;
}

void work(int p,int t){
    int r;
    vector<int> road;

    m=t;mins=INF;
    DFS1(0,p);
    r=root;
    for(int i=0;i<G[r].size();i++){
        Data j=G[r][i];
        if(del[j.v]) continue;
        rd[j.v].clear();
        DFS2(j.v,r,j.v,j.w);
        ans-=census(rd[j.v]);
        for(int k=0;k<rd[j.v].size();k++) road.push_back(rd[j.v][k]);
    }
    del[r]=1;
    ans+=census(road);
    for(int i=0;i<road.size();i++){
        if(road[i]<=k) ans++;
    }
    for(int i=0;i<G[r].size();i++){
        Data j=G[r][i];
        if(del[j.v]) continue;
        work(j.v,son[j.v]);
    }
}

int main(){
    while(1){
        for(int i=1;i<=n;i++) {G[i].clear();rd[i].clear();}
        memset(son,0,sizeof(son));
        memset(del,0,sizeof(del));
        n=k=ans=m=mins=root=0;
        if(!read()) break;
        work(1,n);
        printf("%d\n",ans);
    }

    return 0;
}