POJ 1741 Tree

POJ 1741 Tree

題目概述:腦書p373

給一個邊帶權的樹。求距離小於K的點對數。

思路:

分治。

需要利用樹重心來做切割，才不會使複雜度退化。

//By SCJ

#include<iostream>

#include<cstdio>

#include<algorithm>

#include<vector>

using namespace std;

#define endl '\n'

#define maxn 10000

#define INF 2000000000

struct Edge{

   int to,cost;

   Edge(int a,int b)

   {

       to=a;cost=b;

   }

};

vector<Edge> G[maxn+1];

typedef pair<int,int> pii;

int n,K,centroid[maxn+1],sz[maxn+1];

int get_size(int v,int p)

{

   int cnt=1;

   for(int i=0;i<G[v].size();++i)

   {

       Edge e=G[v][i];

       if(e.to==p||centroid[e.to]) continue;

       cnt+=get_size(e.to,v);

   }

   return sz[v]=cnt;

}

pii find_centroid(int v,int p,int all)

{

   pii res=make_pair(INF,-1);

   int max_sub=0;

   for(int i=0;i<G[v].size();++i)

   {

       Edge e=G[v][i];

       if(e.to==p||centroid[e.to]) continue;

       res=min(res,find_centroid(e.to,v,all));

       max_sub=max(max_sub,sz[e.to]);

   }

   max_sub=max(max_sub,all-sz[v]);

   return min(res,make_pair(max_sub,v));

}

int ans=0;

void get_dis(int v,int p,int d,vector<int> &ds)

{

   ds.push_back(d);

   for(int i=0;i<G[v].size();++i)

   {

       Edge e=G[v][i];

       if(e.to==p||centroid[e.to]) continue;

       get_dis(e.to,v,d+e.cost,ds);

   }

}

int get_path_num(vector<int> &ds)

{

   int res=0;

   sort(ds.begin(),ds.end());

   for(int i=0,j=ds.size()-1;i<ds.size();++i)

   {

       while(i<=j&&ds[i]+ds[j]>K) j--;

       if(i>j) break;

       res+=j-i;

   }

   return res;

}

int solve(int v)

{

   get_size(v,-1);

   int cen=find_centroid(v,-1,sz[v]).second;

   centroid[cen]=1;

   for(int i=0;i<G[cen].size();++i)

   {

       Edge e=G[cen][i];

       if(centroid[e.to]) continue;

       solve(e.to);

   }

   vector<int> ds;

   ds.push_back(0);

   for(int i=0;i<G[cen].size();++i)

   {

       Edge e=G[cen][i];

       if(centroid[e.to]) continue;

       vector<int> tds;

       get_dis(e.to,-1,e.cost,tds);

       ans-=get_path_num(tds);

       //for(int i=0;i<tds.size();++i) ds.push_back(tds[i]);

       ds.insert(ds.end(),tds.begin(),tds.end());

   }

   ans+=get_path_num(ds);

   centroid[cen]=0;

}

int main()

{

//ios::sync_with_stdio(0);

//cin.tie(0);

   //while(cin>>n>>K,n!=0)

   while(~scanf("%d%d",&n,&K))

   {

       if(n==0) break;

       for(int i=1;i<n;++i)

       {

           int a,b,c;

           //cin>>a>>b>>c;

           scanf("%d%d%d",&a,&b,&c);

           G[a].push_back(Edge(b,c));

           G[b].push_back(Edge(a,c));

       }

       ans=0;

       solve(1);

       //cout<<ans<<endl;

       printf("%d\n",ans);

       for(int i=0;i<=n;++i) G[i].clear();

   }

}