HDU 5016 Mart Master II

  1 #include <bits/stdc++.h>
  2 using namespace std;
  3 using PII = pair<int,int>;
  4 const int maxn = 100010;
  5 const int INF = ~0u>>2;
  6 PII d[3][maxn];
  7 struct arc {
  8     int to,w,next;
  9     arc(int x = 0,int y = 0,int z = -1) {
 10         to = x;
 11         w = y;
 12         next = z;
 13     }
 14 } e[maxn<<1];
 15 int head[maxn],sz[maxn],maxson[maxn],mart[maxn],tot,cnt,n;
 16 int ans[maxn];
 17 void add(int u,int v,int w) {
 18     e[tot] = arc(v,w,head[u]);
 19     head[u] = tot++;
 20 }
 21 queue<int>q;
 22 bool done[maxn];
 23 void spfa() {
 24     memset(done,false,sizeof done);
 25     while(!q.empty()){
 26         int u = q.front();
 27         q.pop();
 28         done[u] = false;
 29         for(int i = head[u]; ~i; i = e[i].next){
 30             PII tmp(d[0][u].first + e[i].w,d[0][u].second);
 31             if(d[0][e[i].to] > tmp){
 32                 d[0][e[i].to] = tmp;
 33                 if(!done[e[i].to]){
 34                     done[e[i].to] = true;
 35                     q.push(e[i].to);
 36                 }
 37             }
 38         }
 39     }
 40 }
 41 int dfs(int u,int fa){
 42     sz[u] = 1;
 43     maxson[u] = 0;
 44     for(int i = head[u]; ~i; i = e[i].next){
 45         if(e[i].to == fa || done[e[i].to]) continue;
 46         dfs(e[i].to,u);
 47         sz[u] += sz[e[i].to];
 48         maxson[u] = max(maxson[u],sz[e[i].to]);
 49     }
 50     return sz[u];
 51 }
 52 int root(const int sum,int u,int fa){
 53     int ret = u;
 54     maxson[u] = max(maxson[u],sum - sz[u]);
 55     for(int i = head[u]; ~i; i = e[i].next){
 56         if(e[i].to == fa || done[e[i].to]) continue;
 57         int x = root(sum,e[i].to,u);
 58         if(maxson[x] < maxson[ret]) ret = x;
 59     }
 60     return ret;
 61 }
 62 void update(int u,int w,int fa){
 63     d[1][cnt] = PII(w,u);
 64     d[2][cnt++] = PII(d[0][u].first - w,d[0][u].second);
 65     for(int i = head[u]; ~i; i = e[i].next){
 66         if(e[i].to == fa || done[e[i].to]) continue;
 67         update(e[i].to,w + e[i].w,u);
 68     }
 69 }
 70 void calc(int u,int w,int sg){
 71     cnt = 0;
 72     update(u,w,0);
 73     sort(d[2],d[2] + cnt);
 74     for(int i = 0; i < cnt; ++i){
 75         if(mart[d[1][i].second]) continue;
 76         auto it = lower_bound(d[2],d[2] + cnt,d[1][i]) - d[2];
 77         ans[d[1][i].second] += (cnt - it)*sg;
 78     }
 79 }
 80 void solve(int u){
 81     int rt = root(dfs(u,0),u,0);
 82     done[rt] = true;
 83     calc(rt,0,1);
 84     for(int i = head[rt]; ~i; i = e[i].next){
 85         if(done[e[i].to]) continue;
 86         calc(e[i].to,e[i].w,-1);
 87     }
 88     for(int i = head[rt]; ~i; i = e[i].next){
 89         if(done[e[i].to]) continue;
 90         solve(e[i].to);
 91     }
 92 }
 93 int main() {
 94     int u,v,w;
 95     while(~scanf("%d",&n)) {
 96         memset(head,-1,sizeof head);
 97         int ret = tot = 0;
 98         for(int i = 1; i < n; ++i) {
 99             scanf("%d%d%d",&u,&v,&w);
100             add(u,v,w);
101             add(v,u,w);
102         }
103         for(int i = 1; i <= n; ++i) {
104             scanf("%d",mart + i);
105             if(mart[i]) {
106                 d[0][i] = PII(0,i);
107                 q.push(i);
108             } else d[0][i] = PII(INF,0);
109             ans[i] = 0;
110         }
111         spfa();
112         solve(1);
113         for(int i = 1; i <= n; ++i)
114             ret = max(ret,ans[i]);
115         printf("%d\n",ret);
116     }
117     return 0;
118 }