【点分治】hdu5016 Mart Master II

  1 #pragma comment(linker, "/STACK:1024000000,1024000000")
  2 #include<cstdio>
  3 #include<algorithm>
  4 #include<cstring>
  5 #include<cstdlib>
  6 using namespace std;
  7 #define MAXN 100001
  8 #define INF 1000000000
  9 typedef pair<int,int> Point;
 10 Point td[MAXN],ds[MAXN],p[MAXN],td2[MAXN],ds2[MAXN];
 11 bool is[MAXN];
 12 int n,K,ans[MAXN],en2,en3;
 13 int v[MAXN<<1],w[MAXN<<1],first[MAXN],next[MAXN<<1],en;
 14 void AddEdge(const int &U,const int &V,const int &W)
 15 {
 16     v[++en]=V;
 17     w[en]=W;
 18     next[en]=first[U];
 19     first[U]=en;
 20 }
 21 bool centroid[MAXN];
 22 int size[MAXN];
 23 int calc_sizes(int U,int Fa)
 24 {
 25     int res=1;
 26     for(int i=first[U];i;i=next[i])
 27       if(v[i]!=Fa&&(!centroid[v[i]]))
 28         res+=calc_sizes(v[i],U);
 29     return size[U]=res;
 30 }
 31 Point calc_centroid(int U,int Fa,int nn)
 32 {
 33     Point res=make_pair(INF,-1);
 34     int sum=1,maxv=0;
 35     for(int i=first[U];i;i=next[i])
 36       if(v[i]!=Fa&&(!centroid[v[i]]))
 37         {
 38           res=min(res,calc_centroid(v[i],U,nn));
 39           maxv=max(maxv,size[v[i]]);
 40           sum+=size[v[i]];
 41         }
 42     maxv=max(maxv,nn-sum);
 43     res=min(res,make_pair(maxv,U));
 44     return res;
 45 }
 46 void calc_dis(int U,int Fa,int d)
 47 {
 48     td[en2++]=make_pair(d,U);
 49     for(int i=first[U];i;i=next[i])
 50       if(v[i]!=Fa&&(!centroid[v[i]]))
 51         calc_dis(v[i],U,d+w[i]);
 52 }
 53 void calc_pairs(Point dis[],Point dis2[],int En,int op)
 54 {
 55     sort(dis2,dis2+En);
 56     for(int i=0;i<En;++i)
 57       ans[dis[i].second]+=(dis2+En-lower_bound(dis2,dis2+En,dis[i]))*op;
 58 }
 59 void solve(int U)
 60 {
 61     calc_sizes(U,-1);
 62     int s=calc_centroid(U,-1,size[U]).second;
 63     centroid[s]=1;
 64     for(int i=first[s];i;i=next[i])
 65       if(!centroid[v[i]])
 66         solve(v[i]);
 67     en3=0;
 68     ds[en3]=make_pair(0,s);
 69     ds2[en3++]=make_pair(p[s].first,p[s].second);
 70     for(int i=first[s];i;i=next[i])
 71       if(!centroid[v[i]])
 72         {
 73           en2=0; calc_dis(v[i],s,w[i]);
 74           for(int i=0;i<en2;++i)
 75             td2[i]=make_pair(p[td[i].second].first-td[i].first,p[td[i].second].second);
 76           calc_pairs(td,td2,en2,-1);
 77           memcpy(ds+en3,td,en2*sizeof(Point));
 78           memcpy(ds2+en3,td2,en2*sizeof(Point));
 79           en3+=en2;
 80         }
 81     calc_pairs(ds,ds2,en3,1);
 82     centroid[s]=0;
 83 }
 84 void init()
 85 {
 86     memset(ans,0,(n+1)*sizeof(int));
 87     memset(first,0,(n+1)*sizeof(int));
 88     en=0;
 89 }
 90 void dfs1(int U,int Fa)
 91 {
 92     for(int i=first[U];i;i=next[i])
 93       if(v[i]!=Fa)
 94         {
 95           p[v[i]]=min(p[v[i]],make_pair(p[U].first+w[i],p[U].second));
 96           dfs1(v[i],U);
 97         }
 98 }
 99 void dfs2(int U,int Fa)
100 {
101     for(int i=first[U];i;i=next[i])
102       if(v[i]!=Fa)
103         {
104           dfs2(v[i],U);
105           p[U]=min(p[U],make_pair(p[v[i]].first+w[i],p[v[i]].second));
106         }
107 }
108 int main()
109 {
110     int a,b,c;
111     while(scanf("%d",&n)!=EOF)
112       {
113           init();
114         for(int i=1;i<n;++i)
115             {
116               scanf("%d%d%d",&a,&b,&c);
117               AddEdge(a,b,c);
118               AddEdge(b,a,c);
119             }
120         for(int i=1;i<=n;++i)
121             {
122               scanf("%d",&is[i]);
123               if(!is[i]) p[i]=make_pair(INF,INF);
124               else p[i]=make_pair(0,i);
125             }
126         dfs2(1,-1);
127         dfs1(1,-1);
128         solve(1);
129         int Ans=0;
130         for(int i=1;i<=n;++i) if(!is[i]) Ans=max(Ans,ans[i]);
131         printf("%d\n",Ans);
132       }
133     return 0;
134 }