potentialized_group_dsu.hpp

SECTIONData Structure INCLUDEnoya/potentialized_group_dsu.hpp

Weighted union-find over a possibly noncommutative group. Store the relative value parent(x)^-1x on every parent edge, so path compression multiplies weights from root to leaf in that order. merge(a,b,w) enforces potential(a)^-1potential(b)=w; reversing a union uses the group inverse, which preserves correctness without assuming commutativity.

Verified by unionfind_with_potential_non_commutative_group.

在一般群上维护元素间的相对关系；适合位移、仿射变换等不能只用数值差表示的并查集约束。

Implementation

View on GitHub

SourceExpanded

#ifndef NOYA_POTENTIALIZED_GROUP_DSU_HPP
#define NOYA_POTENTIALIZED_GROUP_DSU_HPP 1

/// @complexity Time: Amortized O(alpha(n)) group operations per query.
/// Space: O(n).

#include <cassert>
#include <optional>
#include <utility>
#include <vector>

namespace noya {

/// @brief Weighted union-find over a possibly noncommutative group. Store the
/// relative value parent(x)^-1*x on every parent edge, so path compression
/// multiplies weights from root to leaf in that order. merge(a,b,w) enforces
/// potential(a)^-1*potential(b)=w; reversing a union uses the group inverse,
/// which preserves correctness without assuming commutativity.
template <class T, class Multiply, class Inverse>
class potentialized_group_dsu {
public:
  potentialized_group_dsu() = default;

  potentialized_group_dsu(int n, T identity, Multiply multiply,
                          Inverse inverse)
      : identity_(std::move(identity)), multiply_(std::move(multiply)),
        inverse_(std::move(inverse)) {
    build(n);
  }

  void build(int n) {
    assert(n >= 0);
    parent_or_size_.assign(n, -1);
    weight_to_parent_.assign(n, identity_);
  }

  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] =
        multiply_(weight_to_parent_[parent], weight_to_parent_[x]);
    return parent_or_size_[x] = root;
  }

  T potential(int x) {
    leader(x);
    return weight_to_parent_[x];
  }

  bool same(int a, int b) { return leader(a) == leader(b); }

  int size(int x) { return -parent_or_size_[leader(x)]; }

  std::optional<T> difference(int a, int b) {
    T weight_a = potential(a);
    T weight_b = potential(b);
    if (leader(a) != leader(b)) {
      return std::nullopt;
    }
    return multiply_(inverse_(weight_a), weight_b);
  }

  bool merge(int a, int b, const 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 multiply_(inverse_(weight_a), weight_b) == difference_ab;
    }

    T root_b_to_root_a = multiply_(
        multiply_(weight_a, difference_ab), inverse_(weight_b));
    if (-parent_or_size_[root_a] < -parent_or_size_[root_b]) {
      std::swap(root_a, root_b);
      root_b_to_root_a = inverse_(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] = std::move(root_b_to_root_a);
    return true;
  }

private:
  T identity_{};
  Multiply multiply_{};
  Inverse inverse_{};
  std::vector<int> parent_or_size_;
  std::vector<T> weight_to_parent_;
};

} // namespace noya

#endif // NOYA_POTENTIALIZED_GROUP_DSU_HPP

#include <cassert>
#include <optional>
#include <utility>
#include <vector>

/// @complexity Time: Amortized O(alpha(n)) group operations per query.
/// Space: O(n).

namespace noya {

/// @brief Weighted union-find over a possibly noncommutative group. Store the
/// relative value parent(x)^-1*x on every parent edge, so path compression
/// multiplies weights from root to leaf in that order. merge(a,b,w) enforces
/// potential(a)^-1*potential(b)=w; reversing a union uses the group inverse,
/// which preserves correctness without assuming commutativity.
template <class T, class Multiply, class Inverse>
class potentialized_group_dsu {
public:
  potentialized_group_dsu() = default;

  potentialized_group_dsu(int n, T identity, Multiply multiply,
                          Inverse inverse)
      : identity_(std::move(identity)), multiply_(std::move(multiply)),
        inverse_(std::move(inverse)) {
    build(n);
  }

  void build(int n) {
    assert(n >= 0);
    parent_or_size_.assign(n, -1);
    weight_to_parent_.assign(n, identity_);
  }

  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] =
        multiply_(weight_to_parent_[parent], weight_to_parent_[x]);
    return parent_or_size_[x] = root;
  }

  T potential(int x) {
    leader(x);
    return weight_to_parent_[x];
  }

  bool same(int a, int b) { return leader(a) == leader(b); }

  int size(int x) { return -parent_or_size_[leader(x)]; }

  std::optional<T> difference(int a, int b) {
    T weight_a = potential(a);
    T weight_b = potential(b);
    if (leader(a) != leader(b)) {
      return std::nullopt;
    }
    return multiply_(inverse_(weight_a), weight_b);
  }

  bool merge(int a, int b, const 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 multiply_(inverse_(weight_a), weight_b) == difference_ab;
    }

    T root_b_to_root_a = multiply_(
        multiply_(weight_a, difference_ab), inverse_(weight_b));
    if (-parent_or_size_[root_a] < -parent_or_size_[root_b]) {
      std::swap(root_a, root_b);
      root_b_to_root_a = inverse_(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] = std::move(root_b_to_root_a);
    return true;
  }

private:
  T identity_{};
  Multiply multiply_{};
  Inverse inverse_{};
  std::vector<int> parent_or_size_;
  std::vector<T> weight_to_parent_;
};

} // namespace noya