Problem p03450 (python)

Fixed submission:
N,M = map(int,input().split())
info = [list(map(int,input().split())) for i in range(M)]
nodes = [i for i in range(N)]
rank = [0]*N
diff_weight = [0]*N

def root(x):
    if nodes[x] == x:
        return x
    root_x = root(nodes[x])
    diff_weight[x] += diff_weight[nodes[x]]
    nodes[x] = root_x
    return nodes[x]

def weight(x):
    root(x)
    return diff_weight[x]

def unite(x,y,w):
    w += weight(x); w -= weight(y)
    x = root(x); y = root(y)
    if x == y:
        return
    if rank[x] >= rank[y]:
        nodes[y] = x
        diff_weight[y] = w
        if rank[x] == rank[y]:
            rank[x] += 1
    else:
        nodes[x] = y
        diff_weight[x] = -w

def same(x,y):
    return root(x) == root(y)

def diff(x,y):
    return weight(y) - weight(x)

ans = "Yes"

for L,R,D in info:
    L -= 1; R -= 1
    if same(L,R):
        if diff(L,R) != D:
            ans = "No"
            break
    else:
        unite(L,R,D)

print(ans)

Buggy submission (verdict: Runtime Error):
import sys
input = sys.stdin.readline

class Unionfind:
    __slots__ = ["nodaes","diff_weight"]

    def __init__(self, n):
        self.nodes = [-1]*n
        self.diff_weight = [0]*n

    def root(self, x):
        if self.nodes[x] < 0:
            return x
        else:
            root_x = self.root(self.nodes[x])
            self.diff_weight[x] += self.diff_weight[self.nodes[x]]
            self.nodes[x] = root_x
            return root_x
    
    def unite(self, x, y, w=0):
        w += self.weight(x); w -= self.weight(y)
        x = self.root(x); y = self.root(y)
        if x == y:
            return
        rank_x = -self.nodes[x]; rank_y = -self.nodes[y]
        if rank_x < rank_y:
            x, y = y, x; w = -w
        if rank_x == rank_y:
            self.nodes[x] -= 1
        self.nodes[y] = x
        self.diff_weight[y] = w

    def weight(self, x):
        self.root(x)
        return self.diff_weight[x]
    
    def diff(self, x, y):
        return self.weight(y) - self.weight(x)
    
    def same(self, x, y):
        return self.root(x) == self.root(y)
    
    def rank(self, x):
        return -self.nodes[self.root(x)]
    
N,M = map(int,input().split())
uf = Unionfind(N)
ans = "Yes"
for _ in range(M):
    L, R, D = map(int,input().split())
    L -= 1; R -= 1
    uf.unite(L, R, D)   
    if uf.diff(L, R) != D:
        ans = "No"
        break
print(ans)

Buggy submission (verdict: Runtime Error):
N,M = map(int,input().split())
info = [list(map(int,input().split())) for i in range(M)]

def root(x):
    # ノードが負 <=> xは根
    if nodes[x] < 0:
        return x
    else:
        parent = nodes[x]
        # 経路圧縮:再帰的に根を調べてその過程で通った要素は全て根に直接つなぎ直す
        nodes[x] = root(nodes[x])
        # 親の重みの累積和をとる. これがxを根に直接繋いだ時の重みになる
        weight[x] += weight[parent]
        return nodes[x]

def unite(x,y,w):
    root_x = root(x); root_y = root(y)
    #根ノードは木の高さを表す(負数で表していることに注意)
    rank_x = -nodes[root_x]; rank_y = -nodes[root_y]
    if rank_x >= rank_y:
        root_l = root_x; root_s = root_y
        l = x; s = y
    else:
        root_l = root_y; root_s = root_x
        l = y; s = x
        w *= -1
    # 繋いだ分だけ木を高くする
    nodes[root_l] += nodes[root_s]
    # 低い木を高い木に繋げる
    nodes[root_s] = root_l
    # 繋ぎ直す時の重みの変更
    weight[root_s] = w - weight[s] + weight[l]

# 根が同じ->True
def same(x,y):
    return root(x) == root(y)

# xからyへの重み
def diff(x,y):
    return weight[y] - weight[x]

# ノードは-1で初期化, Nは頂点数
nodes = [-1]*N
# ウェイトは0で初期化, Nは頂点数
weight = [0]*N

ans = "Yes"

for L,R,D in info:
    L -= 1; R -= 1
    unite(L,R,D)

for L,R,D in info:
    L -= 1; R -= 1
    if diff(L,R) != D:
        ans = "No"
        break

print(ans)