POJ 1741 - Tree

  1 #include <cstdio>
  2 #include <cstring>
  3 using namespace std;
  4 
  5 
  6 #define FOR(p,i,s,t) for(__typeof(p) i=s; i<t; ++i)
  7 #define REP(t,i,n)    FOR(t,i,0,n)
  8 
  9 #define ECH(it, A) for (__typeof(A.begin()) it=A.begin(); it != A.end(); ++it)
 10 #define RST(x,y) memset(x, y, sizeof(x))
 11 #define RST0(x)    RST(x,0)
 12 
 13 typedef int Vt, Lt;
 14 const __typeof(Vt) MAXV = 10005;
 15 
 16 #define MAXE    ((MAXV<<1) - 2)
 17 
 18 Vt Vefw[MAXE], Veh[MAXV], eptr = 0;
 19 struct Vedge {
 20     Vt t;
 21     Lt l;
 22     Vedge() {}
 23     Vedge(Vt _t): t(_t), l(1) {}
 24     Vedge(Vt _t, Lt _l): t(_t), l(_l) {}
 25     void attach(Vt s) {
 26         extern Vedge Vs[];
 27         Vs[eptr].t = this->t, Vs[eptr].l = this->l;
 28         Vefw[eptr] = Veh[s]; Veh[s] = ++eptr;
 29     }
 30 };
 31 #define addedge(s,t,l) ({Vedge e(t,l); e.attach(s);})
 32 Vedge Vs[MAXE];
 33 Vt gcoref_tot;
 34 char gc_8[MAXV];
 35 Vt gc_maxk[MAXV], gc_sumn[MAXV];
 36 
 37 int gc_root;
 38 
 39 void gcoref(Vt i)
 40 {
 41     char _gc8;
 42     if (!(_gc8 = gc_8[i])) gc_8[i] = -1;    // 遍历去环
 43     Vt sumn = 1, maxk = 0;
 44     for(Vt e = Veh[i]; e; e = Vefw[e]) {
 45         Vt t = Vs[--e].t;
 46         if (!gc_8[t]) {
 47             gcoref(t);
 48             sumn += gc_sumn[t];
 49             if (maxk < gc_sumn[t])
 50                 maxk = gc_sumn[t];
 51         }
 52     }
 53     gc_8[i] = _gc8;                // gc_8还有其他用途
 54     if (gcoref_tot - sumn > maxk) maxk = gcoref_tot - sumn;
 55     gc_sumn[i] = sumn, gc_maxk[i] = maxk;
 56     if (gc_maxk[gc_root] > maxk) gc_root = i;
 57 }
 58 
 59 inline Vt gcore(Vt root)
 60 {
 61     gc_maxk[gc_root = root] = gcoref_tot;
 62     gcoref(root);
 63     return gc_root;
 64 }
 65 
 66 int K;
 67 #define dist gc_sumn
 68 Vt sorted_dist[MAXV], sorted_tot;
 69 
 70 void upddis(Vt s)
 71 {
 72     char _gc8;
 73     if (!(_gc8 = gc_8[s])) gc_8[s] = -1;    // 遍历去环
 74     for(Vt e = Veh[s]; e; e = Vefw[e]) {
 75         Vt t = Vs[--e].t;
 76         if (!gc_8[t]) {
 77             if ((dist[t] = dist[s] + Vs[e].l) <= K) {
 78                 sorted_dist[sorted_tot++] = dist[t];
 79                 upddis(t);
 80             }
 81         }
 82     }
 83     gc_8[s] = _gc8;
 84 }
 85 
 86 #include <algorithm>
 87 
 88 int calcpairs(Lt walked, Vt s)
 89 {
 90     sorted_dist[0] = dist[s] = walked, sorted_tot = 1;
 91     upddis(s);
 92     sort(sorted_dist, sorted_dist + sorted_tot);
 93 
 94     int pairs = 0;
 95     Vt i, j;
 96 
 97     i = 0, j = sorted_tot - 1;
 98     while(i < j)
 99         if (sorted_dist[i] + sorted_dist[j] <= K)
100             pairs += j - i++;
101         else --j;
102     return pairs;
103 }
104 
105 Vt subtree_ntot[MAXV];
106 
107 int solve(Vt s)
108 {
109     Vt root = gcore(s);
110     for(Vt e = Veh[root]; e; e = Vefw[e]) {    // 保护子树的结点数
111         Vt t = Vs[--e].t;
112         subtree_ntot[t] = gc_sumn[t];
113     }
114     int pairs = calcpairs(0, root);
115     gc_8[root] = 2;
116     for(Vt e = Veh[root]; e; e = Vefw[e]) {
117         Vt t = Vs[--e].t;
118         if (!gc_8[t]) {
119             gcoref_tot = subtree_ntot[t];
120             pairs -= calcpairs(Vs[e].l, t);
121             pairs += solve(t);
122         }
123     }
124     return pairs;
125 }
126 
127 Vt N;
128 
129 int main(void)
130 {
131 //    freopen("poj1741.txt", "r", stdin);
132     while(scanf("%d%d", &N,&K), N && K) {
133         Vt root = 0;
134         REP(int, i, N-1) {
135             Lt l;
136             Vt s,t;
137             scanf("%d%d%d", &s, &t, &l); --s, --t;
138             addedge(s,t,l), addedge(t,s,l);
139             if (root == t) root = s;
140         }
141         gcoref_tot = N;
142         printf("%d\n", solve(root));
143         RST0(gc_8), RST0(Veh), eptr = 0;
144     }
145     return 0;
146 }