Dsu weighted (#192)

* first draft of dsu weighted

* [auto-verifier] verify commit d124b1f

* golf

* [auto-verifier] verify commit 527432c

* add aizu test now

* [auto-verifier] verify commit f708280

* add docs now

* [auto-verifier] verify commit 82effb2

* do it this way actually

* [auto-verifier] verify commit 73cafbd

* fix bug in test actually

* [auto-verifier] verify commit 035c64c

* fix+golf

* [auto-verifier] verify commit b60cf61

* more golf

* fix

* another fix

* another fix

* [auto-verifier] verify commit 525fcf3

* one last optimization

* [auto-verifier] verify commit bcd04c6

* one last fix

---------

Co-authored-by: GitHub <noreply@github.com>

Author: lrvideckis

.verify-helper/timestamps.remote.json

@@ -43,7 +43,9 @@
 "tests/library_checker_aizu_tests/data_structures/simple_tree_line.test.cpp": "2026-03-01 19:36:27 -0700",
 "tests/library_checker_aizu_tests/data_structures/simple_tree_walk.test.cpp": "2026-03-01 19:36:27 -0700",
 "tests/library_checker_aizu_tests/dsu/dsu.test.cpp": "2026-02-27 15:26:53 -0700",
-"tests/library_checker_aizu_tests/dsu/dsu_bipartite.test.cpp": "2026-02-27 15:26:53 -0700",
+"tests/library_checker_aizu_tests/dsu/dsu_bipartite.test.cpp": "2026-03-06 15:51:34 -0700",
+"tests/library_checker_aizu_tests/dsu/dsu_weighted_aizu.test.cpp": "2026-03-06 15:57:01 -0700",
+"tests/library_checker_aizu_tests/dsu/dsu_weighted_lib_checker.test.cpp": "2026-03-06 15:57:01 -0700",
 "tests/library_checker_aizu_tests/dsu/kruskal_tree_aizu.test.cpp": "2026-02-27 15:26:53 -0700",
 "tests/library_checker_aizu_tests/dsu/line_tree_aizu.test.cpp": "2026-02-27 15:26:53 -0700",
 "tests/library_checker_aizu_tests/dsu/line_tree_lib_checker.test.cpp": "2026-02-27 15:26:53 -0700",

library/dsu/dsu_bipartite.hpp

@@ -8,23 +8,21 @@ struct dsu_bipartite {
   int f(int v) {
     if (p[v] < 0) return v;
     int root = f(p[v]);
-    parity[v] ^= parity[p[v]];
-    return p[v] = root;
+    return parity[v] ^= parity[p[v]], p[v] = root;
   }
+  int size(int v) { return -p[f(v)]; }
+  bool is_bipartite(int v) { return is_bi[f(v)]; }
   bool join(int u, int v) {
-    int root_u = f(u), root_v = f(v);
-    int new_parity = parity[v] ^ parity[u];
-    u = root_u, v = root_v;
-    if (u == v) {
-      is_bi[u] &= new_parity;
+    int x = f(u), y = f(v);
+    int new_parity = parity[u] ^ parity[v];
+    if (x == y) {
+      is_bi[x] &= new_parity;
       return 0;
     }
-    if (p[u] > p[v]) swap(u, v);
-    is_bi[u] &= is_bi[v];
-    parity[v] = new_parity ^ 1;
-    p[u] += p[v], p[v] = u;
+    if (p[x] > p[y]) swap(x, y);
+    is_bi[x] &= is_bi[y];
+    parity[y] = new_parity ^ 1;
+    p[x] += p[y], p[y] = x;
     return 1;
   }
-  int size(int v) { return -p[f(v)]; }
-  bool is_bipartite(int v) { return is_bi[f(v)]; }
 };

library/dsu/dsu_weighted.hpp

@@ -0,0 +1,31 @@
+#pragma once
+//! @code
+//!   dsu_weighted dsu(n);
+//!   dsu.join(u, v, w);
+//!   // we now know a[u] = a[v] + w
+//!   ll w = dsu.diff(u, v);
+//!   // satisfies a[u] = a[v] + w based on prior joins
+//! @endcode
+//! @time O(n + q * \alpha(n))
+//! @space O(n)
+struct dsu_weighted {
+  vi p;
+  vector<ll> d;
+  dsu_weighted(int n): p(n, -1), d(n) {}
+  int f(int u) {
+    if (p[u] < 0) return u;
+    int root = f(p[u]);
+    return d[u] += d[p[u]], 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) {
+    int x = f(u), y = f(v);
+    if (x == y) return 0;
+    w += d[u] - d[v];
+    if (p[x] > p[y]) swap(x, y), w = -w;
+    return p[x] += p[y], p[y] = x, d[y] = w, 1;
+  }
+};

tests/.config/.cppcheck_suppression_list

@@ -43,6 +43,7 @@ knownConditionTrueFalse:../library/strings/suffix_array/suffix_array.hpp:62
 knownConditionTrueFalse:../library/strings/suffix_array/suffix_array_short.hpp:34
 knownConditionTrueFalse:../library/dsu/kruskal_tree.hpp:13
 knownConditionTrueFalse:../library/dsu/dsu.hpp:11
+knownConditionTrueFalse:../library/dsu/dsu_weighted.hpp:29
 constVariable:../kactl/content/numerical/NumberTheoreticTransform.h:30
 constVariable:../kactl/content/graph/CompressTree.h:20
 constVariableReference:../kactl/content/graph/CompressTree.h:20

tests/library_checker_aizu_tests/dsu/dsu_weighted_aizu.test.cpp

@@ -0,0 +1,31 @@
+#define PROBLEM \
+  "https://onlinejudge.u-aizu.ac.jp/problems/DSL_1_B"
+#include "../template.hpp"
+#include "../../../library/dsu/dsu_weighted.hpp"
+#include "../../../library/dsu/dsu.hpp"
+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) {
+      int w;
+      cin >> w;
+      dsu_w.join(u, v, w);
+      dsu.join(u, v);
+      assert(dsu.size(u) == dsu_w.size(u));
+      assert(dsu.size(v) == dsu_w.size(v));
+    } else {
+      ll curr_diff = dsu_w.diff(u, v);
+      if (curr_diff == ll(1e18)) cout << "?\n";
+      else cout << dsu_w.diff(u, v) << '\n';
+      assert(dsu.size(u) == dsu_w.size(u));
+      assert(dsu.size(v) == dsu_w.size(v));
+    }
+  }
+  return 0;
+}

tests/library_checker_aizu_tests/dsu/dsu_weighted_lib_checker.test.cpp

@@ -0,0 +1,36 @@
+#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) {
+      int w;
+      cin >> w;
+      bool joined = dsu_w.join(u, v, w);
+      if (joined) cout << 1 << '\n';
+      else
+        cout << ((dsu_w.diff(u, v) % 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 {
+      ll curr_diff = dsu_w.diff(u, v);
+      if (curr_diff == ll(1e18)) cout << -1 << '\n';
+      else cout << (curr_diff % mod + mod) % mod << '\n';
+      assert(dsu.size(u) == dsu_w.size(u));
+      assert(dsu.size(v) == dsu_w.size(v));
+    }
+  }
+  return 0;
+}

tests/scripts/compile_commented_snippets.sh

@@ -25,7 +25,7 @@ git submodule update
 	echo "vi rhs;"
 	echo "vector<vi> mat;"
 	echo "vector<vector<bool>> grid;"
-	echo "int n,m,k,tl,tr,l,r,l1,r1,l2,r2,s_l,s_r,root_l,root_r,source,sink,total_flow,bccid,u,v,lsz,rsz,cols,cap,num,x,y,i,j,i1,i2,j1,j2,len;"
+	echo "int n,m,k,tl,tr,l,r,l1,r1,l2,r2,s_l,s_r,root_l,root_r,source,sink,total_flow,bccid,u,v,w,lsz,rsz,cols,cap,num,x,y,i,j,i1,i2,j1,j2,len;"
 } >entire_library_without_main
 
 {