как ускорить код python? MST Графы Минимальное остовое дерево

Не могу ускорить код, лимит 1 секунда, задание на codeforces 1242B 0-1 MST

вот то что я написал:

class UF:
    def __init__(self, n):
        self.id = list(range(n))
        self.sz = [1 for i in range(n)]
        self.w = [0 for i in range(n)]

    def find(self, p):
        while p != self.id[p]:
            p = self.id[p]
        return p

    def union(self, a, b, w):
        aRadix = self.find(a)
        bRadix = self.find(b)

        if bRadix != aRadix:
            sza = self.sz[aRadix]
            szb = self.sz[bRadix]
            if sza < szb:
                self.id[aRadix] = bRadix
                self.sz[bRadix] += sza
                self.w[bRadix] += w + self.w[aRadix]
            else:
                self.id[bRadix] = aRadix
                self.sz[aRadix] += szb
                self.w[aRadix] += w + self.w[bRadix]


n, m = map(int, input().split())
uf = UF(n)

conn = [list(map(lambda x: int(x) - 1, input().split())) for i in range(m)]

for i in range(n):
    for j in range(i + 1, n):
        if [i, j] not in conn and [j, i] not in conn:
            uf.union(i, j, 0)

for i, j in conn:
    uf.union(i, j, 1)

print(max(uf.w))

Помогите пожалуйста ускорить!


Ответы (0 шт):