nits/golf

Author: lrvideckis

library/dsu/range_parallel_dsu.hpp

@@ -7,22 +7,17 @@
 //! @space O(n log n)
 struct rp_dsu {
   vector<DSU> dsus;
-  rp_dsu(int n): dsus(bit_width(unsigned(n)), DSU(n)) {}
+  rp_dsu(int n): dsus(bit_width(n + 0u), DSU(n)) {}
   void join(int u, int v, int len, const auto& f) {
     int i = __lg(len);
     join(u, v, f, i);
     join(u + len - (1 << i), v + len - (1 << i), f, i);
   }
   void join(int u, int v, const auto& f, int i) {
-    if (i == 0) {
-      u = dsus[0].f(u);
-      v = dsus[0].f(v);
-      if (!dsus[0].join(v, u)) return;
-      int w = dsus[0].f(u);
-      return f(w, u ^ v ^ w);
-    }
     if (!dsus[i].join(u, v)) return;
-    join(u, v, f, i - 1);
-    join(u + (1 << (i - 1)), v + (1 << (i - 1)), f, i - 1);
+    if (i == 0) return f(u, v);
+    i--;
+    join(u, v, f, i);
+    join(u + (1 << i), v + (1 << i), f, i);
   }
 };

tests/library_checker_aizu_tests/dsu/range_parallel_dsu.test.cpp

@@ -9,20 +9,26 @@ int main() {
   cin.tie(0)->sync_with_stdio(0);
   int n, q;
   cin >> n >> q;
-  rp_dsu dsu(n);
+  rp_dsu r_dsu(n);
   vi y(n);
   for (int i = 0; i < n; i++) cin >> y[i];
   vi x = y;
   int ans = 0;
+  DSU dsu(n);
   auto f = [&](int u, int v) {
-    ans = (ans + 1LL * x[u] * x[v]) % mod;
-    x[u] = (x[u] + x[v]) % mod;
+    u = dsu.f(u);
+    v = dsu.f(v);
+    assert(dsu.join(u, v));
+    int root = dsu.f(u);
+    int other = root ^ u ^ v;
+    ans = (ans + 1LL * x[root] * x[other]) % mod;
+    x[root] = (x[root] + x[other]) % mod;
   };
   vector<vector<pii>> joins(n + 1);
   for (int qq = 0; qq < q; qq++) {
     int k, a, b;
     cin >> k >> a >> b;
-    if (k) dsu.join(a, b, k, f);
+    if (k) r_dsu.join(a, b, k, f);
     joins[k].emplace_back(a, b);
     cout << ans << '\n';
     if (qq == 0 || qq == 1 || qq == 10 || qq == 1000 ||