BZOJ3697 采药人的路径

  1 /**************************************************************
  2     Problem: 3697
  3     User: rausen
  4     Language: C++
  5     Result: Accepted
  6     Time:1752 ms
  7     Memory:15924 kb
  8 ****************************************************************/
  9  
 10 #include <cstdio>
 11 #include <algorithm>
 12  
 13 using namespace std;
 14 const int N = 100005;
 15  
 16 struct edge {
 17     int next, to, v;
 18     edge() {}
 19     edge(int _n, int _t, int _v) : next(_n), to(_t), v(_v) {}
 20 } e[N << 1];
 21   
 22 int first[N], tot;
 23   
 24 struct tree_node {
 25     int sz, dep, dis;
 26     bool vis;
 27 } tr[N];
 28   
 29 int n, k;
 30 int cnt[N << 1];
 31 long long f[N << 1][2], g[N << 1][2], ans;
 32 int Root, Maxsz;
 33 int mx_dep, mx;
 34  
 35 inline int read() {
 36     int x = 0;
 37     char ch = getchar();
 38     while (ch < ‘0‘ || ‘9‘ < ch)
 39         ch = getchar();
 40     while (‘0‘ <= ch && ch <= ‘9‘) {
 41         x = x * 10 + ch - ‘0‘;
 42         ch = getchar();
 43     }
 44     return x;
 45 }
 46  
 47 void Add_Edges(int x, int y, int z) {
 48     e[++tot] = edge(first[x], y, z), first[x] = tot;
 49     e[++tot] = edge(first[y], x, z), first[y] = tot;
 50 }
 51   
 52 void dfs(int p, int fa, int sz) {
 53     int x, y, maxsz = 0;
 54     tr[p].sz = 1;
 55     for (x = first[p]; x; x = e[x].next)
 56         if ((y = e[x].to) != fa && !tr[y].vis) {
 57             dfs(y, p, sz);
 58             tr[p].sz += tr[y].sz;
 59             maxsz = max(maxsz, tr[y].sz);
 60         }
 61     maxsz = max(maxsz, sz - tr[p].sz);
 62     if (maxsz < Maxsz)
 63         Root = p, Maxsz = maxsz;
 64 }
 65   
 66 int get_root(int p, int sz) {
 67     Maxsz = N << 1;
 68     dfs(p, 0, sz);
 69     return Root;
 70 }
 71  
 72 void Dfs(int p, int fa) {
 73     int x, y;
 74     mx_dep = max(mx_dep, tr[p].dep);
 75     if (cnt[tr[p].dis]) ++f[tr[p].dis][1];
 76     else ++f[tr[p].dis][0];
 77     ++cnt[tr[p].dis];
 78     for (x = first[p]; x; x = e[x].next)
 79         if ((y = e[x].to) != fa && !tr[y].vis) {
 80             tr[y].dep = tr[p].dep + 1, tr[y].dis = tr[p].dis + e[x].v;
 81             Dfs(y, p);
 82         }
 83     --cnt[tr[p].dis];
 84 }
 85  
 86 void cal(int p) {
 87     int x, y, j;
 88     mx = 0, g[n][0] = 1;
 89     for (x = first[p]; x; x = e[x].next)
 90         if (!tr[y = e[x].to].vis) {
 91             tr[y].dis = n + e[x].v, tr[y].dep = 1, mx_dep = 1;
 92             Dfs(y, p);
 93             mx = max(mx, mx_dep);
 94             ans += (g[n][0] - 1) * f[n][0];
 95             for (j = -mx_dep; j <= mx_dep; ++j)
 96                 ans += g[n - j][1] * f[n + j][1] + g[n - j][0] * f[n + j][1] + g[n - j][1] * f[n + j][0];           
 97             for (j = n - mx_dep; j <= n + mx_dep; ++j) {
 98                 g[j][0] += f[j][0], g[j][1] += f[j][1];
 99                 f[j][0] = f[j][1] = 0;
100             }
101         }
102     for (j = n - mx; j <= n + mx; ++j)
103         g[j][0] = g[j][1] = 0;
104 }
105  
106 void work(int p, int sz) {
107     int root = get_root(p, sz), x, y;
108     tr[root].vis = 1;
109     cal(root);
110     for (x = first[root]; x; x = e[x].next)
111         if (!tr[y = e[x].to].vis)
112             work(y, tr[y].sz);
113 }
114  
115 int main() {
116     int i, x, y, z;
117     n = read();
118     for (i = 1; i < n; ++i) {
119         x = read(), y = read(), z = read();
120         Add_Edges(x, y, z ? 1 : -1);
121     }
122     work(1, n);
123     printf("%lld\n", ans);
124     return 0;
125 }