Solução do P2 - 2016

p2.py

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

def merge_root(root, i, j):
    good, bad = sorted([root[i], root[j]])
    bad_keys = [k for k, v in root.iteritems() if v == bad]
    for key in bad_keys:
        root[key] = good

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()

    root = {i: i for i in range(n)}

    selected = []
    for e in edges:
        d, i, j = e
        if len(selected) == n-1:
            break
        if root[i] != root[j]:
            merge_root(root, i, j)
            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