bzoj 3697

  1 /**************************************************************
  2     Problem: 3697
  3     User: idy002
  4     Language: C++
  5     Result: Accepted
  6     Time:2480 ms
  7     Memory:14012 kb
  8 ****************************************************************/
  9  
 10 #include <cstdio>
 11 #define N 100010
 12 #define fprintf(...)
 13  
 14 int n;
 15 int head[N], dest[N+N], wght[N+N], next[N+N], ntot;
 16 int fat[N], vis[N], siz[N], bac[N], dis[N], cur_root, cur_siz;
 17 int vv[N+N], svv[N+N][2], cvv[N+N][2];
 18 int stkc[N+N], tpc, stks[N+N], tps;
 19 long long ans;
 20  
 21 void insert( int u, int v, int w ) {
 22     ntot++;
 23     next[ntot] = head[u];
 24     wght[ntot] = w ? 1 : -1;
 25     dest[ntot] = v;
 26     head[u] = ntot;
 27 }
 28 void dfs_root( int u ) {
 29     siz[u] = 1, bac[u] = 0;
 30     for( int t=head[u]; t; t=next[t] ) {
 31         int v=dest[t];
 32         if( vis[v] || v==fat[u] ) continue;
 33         fat[v] = u;
 34         dfs_root(v);
 35         siz[u] += siz[v];
 36         if( siz[v]>bac[u] ) bac[u]=siz[v];
 37     }
 38     if( cur_siz-siz[u]>bac[u] ) bac[u]=cur_siz-siz[u];
 39     if( bac[u]<bac[cur_root] ) cur_root=u;
 40 }
 41 void dfs_dis( int u ) {
 42     if( dis[u]==0 && vv[N+0]>1 )
 43         ans++;
 44     bool ok = vv[N+dis[u]];
 45     if( ok ) {
 46         ans += svv[N-dis[u]][0]+svv[N-dis[u]][1];
 47         cvv[N+dis[u]][1]++;
 48     } else {
 49         ans += svv[N-dis[u]][1];
 50         cvv[N+dis[u]][0]++;
 51     }
 52     stkc[++tpc] = dis[u];
 53     fprintf( stderr, "Arrived %d cur_ans = %lld\n", u, ans );
 54     vv[N+dis[u]]++;
 55     for( int t=head[u]; t; t=next[t] ) {
 56         int v=dest[t], w=wght[t];
 57         if( vis[v] || v==fat[u] ) continue;
 58         dis[v] = dis[u]+w;
 59         fat[v] = u;
 60         dfs_dis( v );
 61     }
 62     vv[N+dis[u]]--;
 63 }
 64 void build_vdcp( int rt ) {
 65     /*   found the core   */
 66     cur_siz = siz[rt];
 67     cur_root = 0;
 68     fat[rt] = rt;
 69     dfs_root( rt );
 70     /*   collect the info  */
 71     rt = cur_root;
 72     vis[rt] = true;
 73     fprintf( stderr, "cur_root = %d\n", rt );
 74     for( int t=head[rt]; t; t=next[t] ) {
 75         int u=dest[t], w=wght[t];
 76         if( vis[u] ) continue;
 77         tpc = 0;
 78         vv[N+0]++;
 79         fat[u] = rt;
 80         dis[u] = w;
 81         dfs_dis( u );
 82         vv[N+0]--;
 83         while( tpc ) {
 84             int v = stkc[tpc--];
 85             if( !cvv[N+v][0] && !cvv[N+v][1] ) continue;
 86             svv[N+v][0] += cvv[N+v][0];
 87             svv[N+v][1] += cvv[N+v][1];
 88             stks[++tps] = v;
 89             cvv[N+v][0] = cvv[N+v][1] = 0;
 90         }
 91     }
 92     while( tps ) {
 93         int v = stks[tps--];
 94         svv[N+v][0] = svv[N+v][1] = 0;
 95     }
 96     for( int t=head[rt]; t; t=next[t] ) {
 97         int u=dest[t];
 98         if( vis[u] ) continue;
 99         build_vdcp( u );
100     }
101 }
102 int main() {
103     scanf( "%d", &n );
104     for( int i=1,u,v,w; i<n; i++ ) {
105         scanf( "%d%d%d", &u, &v, &w );
106         insert( u, v, w );
107         insert( v, u, w );
108     }
109     bac[0] = n;
110     siz[1] = n;
111     build_vdcp(1);
112     printf( "%lld\n", ans );
113 }