MST - Kruskal

gistfile1.py

def distance(points, i, j):
    x1, y1 = points[i]
    x2, y2 = points[j]
    return ((x1 - x2) ** 2 + (y1 - y2) ** 2) ** 0.5

def find_root(height_and_parent, i):
    hi, pi = height_and_parent[i]
    if pi is None:
        return i
    return find_root(height_and_parent, pi)

def merge_subtrees(height_and_parent, ri, rj):
    hri = height_and_parent[ri][0]
    hrj = height_and_parent[rj][0]
    if hri <= hrj:
        height_and_parent[ri] = (hri, rj)
        height_and_parent[rj] = (max(hrj, hri+1), None)
    else:
        height_and_parent[rj] = (hrj, ri)
        height_and_parent[ri] = (max(hri, hrj+1), None)

def soma_conexoes(instance):
    points = [(x, y) for (name, x, y) in instance]
    n = len(points)

    edges = []
    for i in range(n):
        for j in range(i+1, n):
            edges.append((distance(points, i, j), i, j))
    edges.sort()

    height_and_parent = [(0, None) for i in range(n)]

    selected = []
    for e in edges:
        d, i, j = e
        if len(selected) == n-1:
            break
        ri = find_root(height_and_parent, i)
        rj = find_root(height_and_parent, j)
        if ri != rj:
            merge_subtrees(height_and_parent, ri, rj)
            selected.append(e)

    return sum(d for (d, i, j) in selected)


instances = [
    [
        ("Alo", 0.0, 1.0),
        ("C1",  0.0, 0.0),
        ("C2",  1.0, 0.0),
    ],
    [
        ("Alo", 0.0, 1.0),
        ("C1",  50, 20),
        ("C2",  35, -200),
        ("C3",  72.6, 20),
        ("C4",  10, 10),
        ("C5",  50, 21),
        ("C6",  50, 22),
        ("C7",  50, 23),
        ("C8",  50, 24),
        ("C9",  50, 25),
        ("C10",  50, 26),
        ("C11",  50, 27),
    ],
]

answers = [2.0, 288.309188635]

assert abs(soma_conexoes(instances[0]) - answers[0]) < 1e-9
assert abs(soma_conexoes(instances[1]) - answers[1]) < 1e-9