Alex_tz307

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

Adblocker detected! Please consider disabling it...

We've detected AdBlock Plus or some other adblocking software preventing Pastebin.com from fully loading.

We don't have any obnoxious sound, or popup ads, we actively block these annoying types of ads!

Please add Pastebin.com to your ad blocker whitelist or disable your adblocking software.

×