Kruskal's algorithm

kruskal.py

"""
Kruskal's algorithm

Pure Python implementation by Danilo J. S. Bellini
"""
import pytest


def kruskal_sets(graph):
    nodes = frozenset.union(*graph)
    sets = {node: {node} for node in nodes}
    for edge in sorted(graph, key=graph.__getitem__):
        set1, set2 = map(sets.__getitem__, edge)
        if set1 is not set2:
            set_union = set1.union(set2)
            for node in set_union:
                sets[node] = set_union
            yield edge


def kruskal_union_find(graph):
    nodes = frozenset.union(*graph)
    parent = {node: node for node in nodes}
    rank = {node: 0 for node in nodes}
    for edge in sorted(graph, key=graph.__getitem__):
        node1, node2 = edge
        while parent[node1] != node1: # Path compression optimization
            node1 = parent[node1] = parent[parent[node1]]
        while parent[node2] != node2:
            node2 = parent[node2] = parent[parent[node2]]
        root1, root2 = sorted([node1, node2], key=rank.__getitem__)
        if root1 != root2:
            parent[root1] = root2
            if rank[root1] == rank[root2]: # Weighted union optimization
                rank[root2] += 1
            yield edge


def kruskal_networkx(graph):
    import networkx as nx
    nxgraph = nx.Graph(tuple(edge) + ({"weight": weight},)
                       for edge, weight in graph.items())
    for node1, node2, data in nx.minimum_spanning_edges(nxgraph):
        yield frozenset((node1, node2))


@pytest.mark.parametrize("kruskal", [kruskal_sets, kruskal_union_find,
                                     kruskal_networkx])
def test_kruskal(kruskal):
    graph = {frozenset(("A", "B")): 5,
             frozenset(("A", "C")): 4,
             frozenset(("A", "D")): 3,
             frozenset(("A", "E")): 2,
             frozenset(("B", "C")): 1,
             frozenset(("B", "D")): 2,
             frozenset(("B", "E")): 4,
             frozenset(("C", "D")): 3,
             frozenset(("C", "E")): 5,
             frozenset(("D", "E")): 4}
    expected = {frozenset(("B", "C")),
                frozenset(("A", "E")),
                frozenset(("B", "D")),
                frozenset(("A", "D"))}
    assert set(kruskal(graph)) == expected