带权/扩展域并查集

What · 是什么

带权并查集（weighted union-find）在普通并查集的每个节点上额外记录一个『到父节点的边权』weight[x]，使每条树边承载任意可结合代数结构（模 3 加法群、异或、普通整数距离等）。经典应用：食物链式『种类关系』（模 3 群：x 与 y 同类/吃/被吃），以及扩展域并查集（把每个元素拆成多个状态点，用普通并查集维护状态集合）。二者等效：模数小的种类关系常用扩展域（好写、不易错），一般关系用带权。注意：路径压缩时需沿路径累加边权；按秩/大小合并时需推导新连边的权值。与普通并查集复杂度同为近似 O(1)/操作（路径压缩+启发式合并，反阿克曼 O(α(n))）。

Why · 为什么

How · 实现

伪代码

class WeightedDSU:
  # weight[x]: x 到根(势差), 模 MOD
  def find(x):
    if parent[x] != x:
      r, w = find(parent[x])
      weight[x] = (weight[x] + w) % MOD
      parent[x] = r
    return parent[x], weight[x]
  def diff(x, y):  # 需同根
    return (weight[x] - weight[y]) % MOD
  def merge(x, y, d):  # 声明 x 相对 y 差为 d
    rx, wx = find(x); ry, wy = find(y)
    if rx == ry: return (wx - wy) % MOD == d % MOD  # 冲突?
    # 把 rx 挂到 ry: 新边权 w 使 wx + w - wy = d => w = d + wy - wx
    parent[rx] = ry; weight[rx] = (d + wy - wx) % MOD
    return True

Python 实现

class WeightedDSU:
    """带权并查集：维护每个节点到其根的距离/相对关系.
    weight[x] = 从 x 到 root 的"势差"(自定义运算, 例如模 MOD 意义下的差/异或/距离).
    这里默认支持模 MOD 加法群 (食物链式 种类关系), 也支持普通距离.
    """
    def __init__(self, n, MOD=1):
        self.parent = list(range(n))
        self.weight = [0] * n   # 到根节点的权值(取模 MOD 或普通)
        self.size = [1] * n
        self.MOD = MOD

    def find(self, x):
        # 路径压缩: 递归累加权值
        if self.parent[x] != x:
            r, w = self.find(self.parent[x])
            self.weight[x] = (self.weight[x] + w) % self.MOD
            self.parent[x] = r
        return self.parent[x], self.weight[x]

    def related(self, x, y):
        return self.find(x)[0] == self.find(y)[0]

    def diff(self, x, y):
        """x 相对 y 的差(带权值), 要求 x,y 同根. 返回 weight[x] - weight[y] (mod MOD)"""
        rx, wx = self.find(x)
        ry, wy = self.find(y)
        assert rx == ry
        return (wx - wy) % self.MOD

    def merge(self, x, y, d):
        """声明: x 相对 y 的差为 d, 即 weight[x] - weight[y] == d (mod MOD).
        若 x,y 已同根则返回 False(冲突), 否则合并返回 True."""
        rx, wx = self.find(x)
        ry, wy = self.find(y)
        if rx == ry:
            return (wx - wy) % self.MOD == d % self.MOD
        # 把 rx 的根接到 ry 的根.  需要: weight[rx] = wx' 使 wx'+wx - wy == d
        # => 从 x 走: weight[x]=wx 相对 rx, 若令 rx 相对 ry 为 w, 则 x 相对 y = wx + w - wy = d
        # => w = d + wy - wx
        w = (d + wy - wx) % self.MOD
        if self.size[rx] < self.size[ry]:
            self.parent[rx] = ry
            self.weight[rx] = w
            self.size[ry] += self.size[rx]
        else:
            self.parent[ry] = rx
            # 此时 ry 接到 rx: 需要 y 相对 rx 的权 = wy', 使 wx - (wy'+wy)=d => wy' = wx - wy - d
            self.weight[ry] = (wx - wy - d) % self.MOD
            self.size[rx] += self.size[ry]
        return True


if __name__ == "__main__":
    # 测试 1: 模 3 食物链式关系 (P2024 风格)
    dsu = WeightedDSU(6, MOD=3)
    # 1 吃 2 (x相对y差为1), 2 吃 3, 则 1 吃 3 (差2 = -1)
    assert dsu.merge(0, 1, 1)
    assert dsu.merge(1, 2, 1)
    assert dsu.diff(0, 2) == 2   # 1 吃 3
    # 3 与 1 同类? 1-3差=2 ≠0, 冲突
    assert dsu.diff(0, 2) != 0
    # 声明 3 与 0 同类 -> 冲突(已经 0 吃 3)
    assert not dsu.merge(0, 2, 0)
    # 测试 2: 普通距离 (银河英雄传说式, MOD=1 无取模)
    d2 = WeightedDSU(4, MOD=1 << 60)
    assert d2.merge(0, 1, 5)      # weight[0]-weight[1]=5
    assert d2.diff(0, 1) == 5
    assert d2.merge(2, 0, 3)      # weight[2]-weight[0]=3
    assert d2.diff(2, 1) == 8
    # 测试 3: 随机与暴力交叉验证
    import random
    random.seed(42)
    for trial in range(400):
        n = random.randint(2, 10)
        d = WeightedDSU(n, MOD=1009)
        val = [random.randrange(1009) for _ in range(n)]
        for step in range(150):
            x, y = random.randrange(n), random.randrange(n)
            dd = random.randrange(1009)
            related = d.related(x, y)
            if not related:
                rx = d.find(x)[0]
                comp = [j for j in range(n) if d.find(j)[0] == rx]
            got = d.merge(x, y, dd)
            if related:
                assert got == ((val[x] - val[y]) % 1009 == dd)
            else:
                assert got == True
                delta = (dd - ((val[x] - val[y]) % 1009)) % 1009
                for j in comp:
                    val[j] = (val[j] + delta) % 1009
        for i in range(n):
            for j in range(n):
                if d.related(i, j):
                    assert d.diff(i, j) == (val[i] - val[j]) % 1009, (trial, i, j)
    print("dsu-weighted: all tests passed")

易错点

• 路径压缩时忘记把父节点到根的权累加到自己身上
• 合并时新连边权计算符号反了（要现场从 diff 定义推导，别背公式）
• 模 M 下 diff 取负需 %M 归一化；Python 负数取模与数学意义一致但直接减可能为负
• 扩展域并查集要把每个元素拆成 n*k 个状态并正确映射合并对
• 只按秩/大小合并时，若忽略 weight 的更新会导致后续查询全错
• find 递归过深在极端数据下爆栈（用迭代或调大递归上限）

例题

信源

• OI Wiki✓ https://oi-wiki.org/ds/dsu/
• OI Wiki✓ https://oi-wiki.org/ds/dsu/#带权并查集
• 书籍✓ CPH ch15.2 Disjoint Set Union