< uruneaさんのAI生成コード置き場

重み付きUnion-Find

日時はサーバーが記録したものです。更新すると更新日時が変わるので、不都合がある場合は新しく投稿してください。

import sys
sys.setrecursionlimit(10**6)

class P_UnionFind:
    def __init__(self, n):
        self.n = n
        self.parents = list(range(n))
        self.sizes = [1] * n
        # diff_weight[x] は parent[x] に対する x の重みの差 (weight(x) - weight(parent[x]))
        self.diff_weight = [0] * n

    def find(self, x):
        if self.parents[x] == x:
            return x
        else:
            root = self.find(self.parents[x])
            self.diff_weight[x] += self.diff_weight[self.parents[x]]
            self.parents[x] = root
            return root

    def weight(self, x):
        # 根に対する x の重みを返す
        self.find(x) # 経路圧縮を行って diff_weight を更新
        return self.diff_weight[x]

    def unite(self, x, y, w):
        rx = self.find(x)
        ry = self.find(y)
        if rx == ry:
            # 矛盾しているかどうかのチェックが必要ならここで行う
            # diff(x, y) == w なら矛盾なし
            return False

        if self.sizes[rx] < self.sizes[ry]:
            rx, ry = ry, rx
            w = -w
        
        # 結合処理
        # weight(y) - weight(x) = w
        # (weight(ry) + diff[y]) - (weight(rx) + diff[x]) = w
        # weight(ry) - weight(rx) = w + diff[x] - diff[y]
        self.parents[ry] = rx
        self.diff_weight[ry] = w + self.weight(x) - self.weight(y)
        self.sizes[rx] += self.sizes[ry]
        return True

    def same(self, x, y):
        return self.find(x) == self.find(y)

    def get_size(self, x):
        return self.sizes[self.find(x)]

    def diff(self, x, y):
        if not self.same(x, y):
            return None
        return self.weight(y) - self.weight(x)


#nはノードの数
uf = P_UnionFind(n)