potentialized_dsu.hpp¶
Weighted union-find over an additive group. merge(a, b, w) enforces potential(b) - potential(a) = w.
Verified by unionfind_with_potential.
维护带势能差的并查集，支持加入两点差值约束并查询相对势能；适合加权并查集模型。
Implementation¶
#ifndef NOYA_POTENTIALIZED_DSU_HPP
#define NOYA_POTENTIALIZED_DSU_HPP 1
/// @complexity Time: Amortized O(alpha(n)) find/merge/difference.
/// Space: O(n).
#include <cassert>
#include <optional>
#include <utility>
#include <vector>
namespace noya {
/// @brief Weighted union-find over an additive group.
/// merge(a, b, w) enforces potential(b) - potential(a) = w.
template <class T> struct potentialized_dsu {
std::vector<int> parent_or_size;
std::vector<T> weight_to_parent;
potentialized_dsu() = default;
explicit potentialized_dsu(int n) { build(n); }
/// @brief Reset to n singleton components with zero potentials.
void build(int n) {
assert(n >= 0);
parent_or_size.assign(n, -1);
weight_to_parent.assign(n, T{});
}
/// @brief Return the representative of x with path compression.
int leader(int x) {
assert(0 <= x && x < int(parent_or_size.size()));
if (parent_or_size[x] < 0) {
return x;
}
int parent = parent_or_size[x];
int root = leader(parent);
weight_to_parent[x] += weight_to_parent[parent];
return parent_or_size[x] = root;
}
/// @brief Return potential(x) relative to its component representative.
T potential(int x) {
leader(x);
return weight_to_parent[x];
}
/// @brief Return whether a and b are in the same component.
bool same(int a, int b) { return leader(a) == leader(b); }
/// @brief Return the size of the component containing x.
int size(int x) { return -parent_or_size[leader(x)]; }
/// @brief Return potential(b) - potential(a), or nullopt if disconnected.
std::optional<T> difference(int a, int b) {
if (!same(a, b)) {
return std::nullopt;
}
return potential(b) - potential(a);
}
/// @brief Add a potential constraint; return false only if it contradicts
/// existing constraints.
bool merge(int a, int b, T difference_ab) {
T weight_a = potential(a);
T weight_b = potential(b);
int root_a = leader(a);
int root_b = leader(b);
if (root_a == root_b) {
return weight_b - weight_a == difference_ab;
}
T root_b_to_root_a = difference_ab + weight_a - weight_b;
if (-parent_or_size[root_a] < -parent_or_size[root_b]) {
std::swap(root_a, root_b);
root_b_to_root_a = -root_b_to_root_a;
}
parent_or_size[root_a] += parent_or_size[root_b];
parent_or_size[root_b] = root_a;
weight_to_parent[root_b] = root_b_to_root_a;
return true;
}
};
} // namespace noya
#endif // NOYA_POTENTIALIZED_DSU_HPP
#include <cassert>
#include <optional>
#include <utility>
#include <vector>
/// @complexity Time: Amortized O(alpha(n)) find/merge/difference.
/// Space: O(n).
namespace noya {
/// @brief Weighted union-find over an additive group.
/// merge(a, b, w) enforces potential(b) - potential(a) = w.
template <class T> struct potentialized_dsu {
std::vector<int> parent_or_size;
std::vector<T> weight_to_parent;
potentialized_dsu() = default;
explicit potentialized_dsu(int n) { build(n); }
/// @brief Reset to n singleton components with zero potentials.
void build(int n) {
assert(n >= 0);
parent_or_size.assign(n, -1);
weight_to_parent.assign(n, T{});
}
/// @brief Return the representative of x with path compression.
int leader(int x) {
assert(0 <= x && x < int(parent_or_size.size()));
if (parent_or_size[x] < 0) {
return x;
}
int parent = parent_or_size[x];
int root = leader(parent);
weight_to_parent[x] += weight_to_parent[parent];
return parent_or_size[x] = root;
}
/// @brief Return potential(x) relative to its component representative.
T potential(int x) {
leader(x);
return weight_to_parent[x];
}
/// @brief Return whether a and b are in the same component.
bool same(int a, int b) { return leader(a) == leader(b); }
/// @brief Return the size of the component containing x.
int size(int x) { return -parent_or_size[leader(x)]; }
/// @brief Return potential(b) - potential(a), or nullopt if disconnected.
std::optional<T> difference(int a, int b) {
if (!same(a, b)) {
return std::nullopt;
}
return potential(b) - potential(a);
}
/// @brief Add a potential constraint; return false only if it contradicts
/// existing constraints.
bool merge(int a, int b, T difference_ab) {
T weight_a = potential(a);
T weight_b = potential(b);
int root_a = leader(a);
int root_b = leader(b);
if (root_a == root_b) {
return weight_b - weight_a == difference_ab;
}
T root_b_to_root_a = difference_ab + weight_a - weight_b;
if (-parent_or_size[root_a] < -parent_or_size[root_b]) {
std::swap(root_a, root_b);
root_b_to_root_a = -root_b_to_root_a;
}
parent_or_size[root_a] += parent_or_size[root_b];
parent_or_size[root_b] = root_a;
weight_to_parent[root_b] = root_b_to_root_a;
return true;
}
};
} // namespace noya