ポテンシャル付き Union Find ¹ と呼ばれるものの一般化について、いろんなパターンを網羅しようとすると混乱したので、考えた結果をメモします。

前提

$R$ を可換整域 ² 、その商体を $\mathrm{Frac}(R)$ とします。以下、四則演算・逆元・単元（ $1$ ）などはすべて $\mathrm{Frac}(R)$ 上で考えるものとします。

$S \subset R$ を乗法群 $R^{×}$ の部分群とします ³。すなわち次の条件を満たすとします。

$1 \in S$
$\forall a \in S$ について $a^{-1} \in S$
$\forall a,\ \forall b \in S$ について $ab \in S$

さらに $\overline{R} \subset \mathrm{Frac}(R)$ を

$\overline{R} = R \cup \{(1-a)^{-1}\cdot x\ |\ a \in S,\ a \neq 1,\ x \in R \}$

で定義します。
$\overline R$ は和については閉じていますが、積については閉じていないことに注意してください。ただし $R$ の元をかける操作については閉じています。
$\overline R$ は後述のクエリの結果出てくる可能性がある値と（集合として）一致します。理論上は$\overline R$ の代わりに $\mathrm{Frac}(R)$ で考えるとらくかもしれませんが、実装上 $\mathrm{Frac}(R)$ まで拡張しなくて良い場合もあるのでこの定義にしています ⁴。

クエリ

$i=1,\ \cdots,\ N$ について $v_i \in R$ が対応しているとします。最初 $v_i$ についての情報はありません。
次の $3$ 種類のクエリを処理したいです。以下クエリにおいて $1\le i,\ j\le N$ 、 $a\in S$ 、 $b \in R$ とします。

$1$ $i$ $j$ $a$ $b$
$v_j=a\cdot v_i+b$ という関係式を追加する ⁵

$2$ $i$
$v_i$ の値の可能性がいくつ存在するか ⁶ 、およびひとつの場合はその値を答える

$3$ $i\ j$
$v_i$ および $v_j$ の値が定まっていないとき、その関係式を答える。すなわち $v_j=a\cdot v_i + b$ をみたす $a \in S,\ b \in R$ （関係式がまだない場合はその旨）を答える ⁷ 。

性質・考察

前節のクエリを「関係式付き Union Find クエリ」と呼びます。
このクエリと通常の Union Find クエリ ⁸ と比べてみましょう。
$a$ が可逆であることから、通常の UF の意味で連結な $2$ 点の間には関係式が定まることが分かります。
実はこの逆は成立しません。これは関係式付き Union Find クエリでは、（通常の UF の意味で）「連結」ではない $2$ 点についても、両方の値が分かっている場合があるからです ⁹ 。

連結成分というと、関係式が分かっているという意味と、通常の UF クエリの意味で連結であるという意味のふたつがあることになります。通常の UF クエリの意味での連結成分を「グループ」と呼ぶことにします。ある頂点の値が確定すれば、同じグループ内のすべての頂点の値が確定します。値が確定しているグループを Type A 、そうでないグループを Type B と呼びます。クエリ 1 における新たな情報は、次の $5$ パターンがあります。

追加情報がない
矛盾を発生する
Type B のグループにループを発生することによって、そのグループの値をすべて確定する
Type A のグループと Type B のグループを結ぶ関係式によって、 Type B 側の頂点の値をすべて確定する
$2$ つの異なる Type B のグループについて関係式を与えることによって、ひとつの Type B のグループにする

ここで、ひとつの Type B のグループ内の $2$ 点間についてのクエリでは、 1. 2. 3. のいずれかになることがポイントです。つまり、連結グループ内に追加情報が少しでもあればグループ内のすべての頂点の値を確定できます。なおここで $R$ が整域であることを使っています ¹⁰。

具体的には、 $v_j=a_1 \cdot v_i + b_1$ という関係式がある $2$ 点について、新たに $v_j = a_2 \cdot v_i + b_2$ という情報を追加したとしましょう。 $v_j$ を消去すると $(a_1-a_2)v_i=b_2-b_1$ となります。ここで $a_1=a_2$ かつ $b_1=b_2$ であれば追加情報はありません。そうでなく $a_1=a_2$ であれば矛盾が発生します。そうでなければ $v_i=(b_2-b_1) \cdot (a_1-a_2)^{-1}$ となり $v_i$ が（したがって $v_j$ も）確定します。さらに $a_1-a_2 =a_1(1-a_2\cdot a_1^{-1}) \in S$ かつ $b_2-b_1 \in R$ なので $v_i \in \overline{R}$ も分かります。

実装の方針

通常の Union Find の実装に加えて、各グループでの代表元との関係を持っておきます。すなわち $i$ が属するグループの代表元を $r_i$ とすると $v_i = a\cdot v_{r_i} + b$ となる $(a,\ b)$ の組を持っておけばよいです。さらに値が定まっている場合はその値も持っておきます。値が確定する場合は、グループ内のすべての頂点における値を明示的に求めても良いですが、代表元の値だけ持っておいて必要に応じて計算する方法でも実現できます。

一個値がわかれば代表元の値もわかるので、単に代表元の値がわかってるかどうかと、わかってるならいくつかさえ持っておけば判明したときに全部に適用する必要とかはなさそうです
— のいみ (@noimi_kyopro) December 3, 2022

例

実数体

$R=\overline{R}=\mathbb{R},\ S=\mathbb{R}-\{0\}$ です。
理論的にはきれいですが、 Float 誤差の問題もあり、正確に実現するのは難しそうです。

mod P

$R=\overline{R}=\mathbb{Z}/p\mathbb{Z},\ S=(\mathbb{Z}/p\mathbb{Z})^{×}$ です。こちらは理論も実装も扱いやすいですね。 $P$ が前計算できないぐらい大きいときは逆元の計算に $\log$ が付きます。

Z/2Z

上に含まれますが $P=2$ の場合です。この場合はクエリ $1$ をいくらやっても値が確定することはありません。

整数環（差のみ）

$R=\overline{R}=\mathbb{Z},\ S=\{1\}$ です。関係式では $2$ 数の差が与えられます（つまり $v_j=v_i+b$ の形で表されます） ¹¹ 。

整数環（和・差）

$R=\mathbb{Z},\ S=\{1,-1\},\ \overline{R}=\{\displaystyle\frac{x}{2}\ |\ x\in \mathbb{Z}\} \subsetneqq \mathrm{Frac}\mathbb{(Z)}$ です。先ほどの例の拡張で、 $2$ 数の和または差が与えられます。ここで初めて $\overline{R}$ を定義した効果がでてきています。 $S$ は可逆元のみを動くので、 $v_i$ として出てくる可能性があるのは半整数 ¹² のみです。実装上は、有理数クラスのようなものは不要で、 $2$ 倍の値を整数として保持するので十分です。見た目が気持ち悪いのがやや気になりますが。