# USACO 2013 US Open, Gold Problem 2. Yin and Yang

May 4th, 2021
597
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
1. #include <bits/stdc++.h>
2. #include <ext/pb_ds/assoc_container.hpp>
3. #include <ext/pb_ds/tree_policy.hpp>
4.
5. using namespace std;
6. using namespace __gnu_pbds;
7. using oset = tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update>;
8. using int64 = long long;
9.
10. ifstream fin("yinyang.in");
11. ofstream fout("yinyang.out");
12.
13. const int MAXN = 1e5 + 5;
14. int sz[MAXN];
15. bitset<MAXN> viz;
16. vector<pair<int,int>> G[MAXN];
17.
18. void find_size(int u, int parent) {
19.   sz[u] = 1;
20.   for (auto v : G[u])
21.     if (v.first != parent && !viz[v.first]) {
22.       find_size(v.first, u);
23.       sz[u] += sz[v.first];
24.     }
25. }
26.
27. int find_centroid(int u, int parent, int n) {
28.   for (auto v : G[u])
29.     if (v.first != parent && !viz[v.first] && sz[v.first] > (n >> 1))
30.       return find_centroid(v.first, u, n);
31.   return u;
32. }
33.
34. int offset, cnt[MAXN << 1], paths_rest[MAXN << 1], paths_no_rest[MAXN << 1];
35. vector<pair<int,bool>> paths;
36.
37. void find_paths(int u, int parent, int sum) {
38.   paths.emplace_back(sum, cnt[sum + offset] > 0);
39.   ++cnt[sum + offset];
40.   for (auto v : G[u])
41.     if (v.first != parent && !viz[v.first])
42.       find_paths(v.first, u, sum + v.second);
43.   --cnt[sum + offset];
44. }
45.
46. int64 solve_tree(int u, int n) {
47.   for (int i = offset - n + 1; i <= offset + n - 1; ++i)
48.     paths_rest[i] = paths_no_rest[i] = 0;
49.   int64 ans = 0;
50.   for (auto v : G[u])
51.     if (!viz[v.first]) {
52.       paths.clear();
53.       find_paths(v.first, u, v.second);
54.       for (auto p : paths) {
55.         if (p.second || p.first == 0)
56.           ans += paths_no_rest[offset - p.first];
57.         ans += paths_rest[offset - p.first];
58.         if (p.second && p.first == 0)
59.           ++ans;
60.       }
61.       for (auto p : paths)
62.         if (p.second)
63.           ++paths_rest[p.first + offset];
64.         else ++paths_no_rest[p.first + offset];
65.     }
66.   return ans;
67. }
68.
69. int64 build(int u, int parent) {
70.   find_size(u, 0);
71.   int c = find_centroid(u, 0, sz[u]);
72.   viz[c] = true;
73.   int64 ans = solve_tree(c, sz[u]);
74.   for (auto v : G[c])
75.     if (!viz[v.first])
76.       ans += build(v.first, c);
77.   return ans;
78. }
79.
80. void test_case() {
81.   int N;
82.   fin >> N;
83.   offset = N;
84.   for (int i = 1; i < N; ++i) {
85.     int u, v, w;
86.     fin >> u >> v >> w;
87.     w = (w << 1) - 1;
88.     G[u].emplace_back(v, w);
89.     G[v].emplace_back(u, w);
90.   }
91.   fout << build(1, 0) << '\n';
92. }
93.
94. void solve() {
95.   int T = 1;
96.   for (int tc = 0; tc < T; ++tc)
97.     test_case();
98. }
99.
100. void close_files() {
101.   fin.close();
102.   fout.close();
103. }
104.
105. int main() {
106.   solve();
107.   close_files();
108.   return 0;
109. }
110.
RAW Paste Data