как ускорить код 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))
Помогите пожалуйста ускорить!