Weighted DSU
Disjoint set union where vertices have their own "potential value". The potential value of a vertex \(i\) is denoted as \(w(i)\), and the differences of potential values between vertices in a same group are always well defined.
fn union(&mut self, u: usize, v: usize, dw: i64)
addes a "rule" stating that \(u\) and \(v\) are in a same group, and \(w(u) - w(v) = dw\). Based on posed rules ahead, fn get_pot_diff(&mut self, u: usize, v: usize) -> Option<i64>
returns \(w(u) - w(v)\) if \(u\) and \(v\) are in a same group.
potdiff[i]
in the code is defined as \(w(i) - w(r_i)\), where \(r_i\) is a root of a group \(i\) is in.
Code
struct WeightDSU {
parent: Vec<usize>,
size: Vec<usize>,
num: usize,
potdiff: Vec<i64>, // potdiff[x] == w(x) - w(p)
}
impl WeightDSU {
fn new(n: usize) -> Self {
Self {
parent: (0..n).collect(),
size: vec![1; n],
num: n,
potdiff: vec![0; n],
}
}
// Returns the root of x.
fn find_root(&mut self, mut x: usize) -> usize {
while self.parent[x] != x {
let p = self.parent[x];
self.potdiff[x] += self.potdiff[p];
self.parent[x] = self.parent[p];
x = self.parent[x];
}
x
}
// Returns the root of x, namely xr, with w(x) - w(xr).
fn find_root_with_pdiff(&mut self, mut x: usize) -> (usize, i64) {
let mut pd = 0;
while self.parent[x] != x {
let p = self.parent[x];
self.potdiff[x] += self.potdiff[p];
self.parent[x] = self.parent[p];
pd += self.potdiff[x];
x = self.parent[x];
}
(x, pd)
}
// Unions groups of u and v, with w(u) - w(v) = dw. If u and v are already in a same group,
// if w(u) - w(v) == dv then it returns Some(true), otherwise Some(false). If they weren't
// in a same group, then it returns None and unions two groups following the given dw.
fn union(&mut self, u: usize, v: usize, dw: i64) -> Option<bool> {
let (ur, pu) = self.find_root_with_pdiff(u);
let (vr, pv) = self.find_root_with_pdiff(v);
let nw = dw - pu + pv;
if ur != vr {
self.num -= 1;
if self.size[ur] < self.size[vr] {
self.parent[ur] = vr;
self.size[vr] += self.size[ur];
self.potdiff[ur] = nw;
} else {
self.parent[vr] = ur;
self.size[ur] += self.size[vr];
self.potdiff[vr] = -nw;
}
None
} else {
Some(nw == 0)
}
}
// Returns the size of a group x is in.
fn get_size(&mut self, x: usize) -> usize {
let r = self.find_root(x);
self.size[r]
}
// Returns Some(w(u) - w(v)) if u and v are in the same group, otherwise None.
fn get_pot_diff(&mut self, u: usize, v: usize) -> Option<i64> {
let (ur, pu) = self.find_root_with_pdiff(u);
let (vr, pv) = self.find_root_with_pdiff(v);
if ur == vr {
Some(pu - pv)
} else {
None
}
}
}