2016年9月2日 星期五

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();
   }
}

沒有留言:

張貼留言