first draft of dsu weighted

Author: lrvideckis

library/dsu/dsu_weighted.hpp

@@ -0,0 +1,27 @@
+#pragma once
+struct dsu_weighted {
+  int n;
+  vi p;
+  vector<ll> d;
+  dsu_weighted(int n): n(n), p(n, -1), d(n) {}
+  int f(int u) {
+    if (p[u] < 0) return u;
+    int root = f(p[u]);
+    d[u] += d[p[u]];
+    return p[u] = root;
+  }
+  int size(int u) { return -p[f(u)]; }
+  ll diff(int u, int v) {
+    return f(u) == f(v) ? d[v] - d[u] : 1e18;
+  }
+  bool join(int u, int v, ll w) {
+    w += d[u] - d[v];
+    u = f(u), v = f(v);
+    if (u == v) return 0;
+    if (p[u] > p[v]) swap(u, v), w = -w;
+    p[u] += p[v];
+    p[v] = u;
+    d[v] = w;
+    return 1;
+  }
+};

tests/library_checker_aizu_tests/dsu/dsu_weighted_lib_checker.test.cpp

@@ -0,0 +1,40 @@
+#define PROBLEM \
+  "https://judge.yosupo.jp/problem/unionfind_with_potential"
+#include "../template.hpp"
+#include "../../../library/dsu/dsu_weighted.hpp"
+#include "../../../library/dsu/dsu.hpp"
+const int mod = 998'244'353;
+int main() {
+  cin.tie(0)->sync_with_stdio(0);
+  int n, q;
+  cin >> n >> q;
+  dsu_weighted dsu_w(n);
+  DSU dsu(n);
+  while (q--) {
+    int type, u, v;
+    cin >> type >> u >> v;
+    if (type == 0) {
+      assert(dsu.size(u) == dsu_w.size(u));
+      assert(dsu.size(v) == dsu_w.size(v));
+      int w;
+      cin >> w;
+      ll curr_diff = dsu_w.diff(u, v);
+      if (curr_diff == 1e18) {
+        assert(dsu_w.join(u, v, w));
+        cout << 1 << '\n';
+      } else
+        cout << ((curr_diff % mod + mod) % mod == w)
+             << '\n';
+      dsu.join(u, v);
+      assert(dsu.size(u) == dsu_w.size(u));
+      assert(dsu.size(v) == dsu_w.size(v));
+    } else {
+      assert(dsu.size(u) == dsu_w.size(u));
+      assert(dsu.size(v) == dsu_w.size(v));
+      ll curr_diff = dsu_w.diff(u, v);
+      if (curr_diff == 1e18) cout << -1 << '\n';
+      else cout << (curr_diff % mod + mod) % mod << '\n';
+    }
+  }
+  return 0;
+}