Demonstration of NetworkX (minimum_spanning_tree)

mst.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""mst.py: demonstrate NetworkX (minimum_spanning_tree)
"""
import networkx as nx
import numpy as np

def main():
    G = setup_graph(5)
    dump_spanning(nx.minimum_spanning_tree(G))

def setup_graph(step=5):
    G = nx.Graph()

    # grid network step x step
    node_indices = np.ndindex(step, step)
    for i in node_indices:
        G.add_node(
            i[0] * step + i[1],
            coords=np.array(i))

    order = len(G)
    locations = nx.get_node_attributes(G, 'coords')
    for n in G.nodes_iter():
        if (n + 1) % step != 0:
            G.add_edge(
                n, n + 1,
                weight=distance(locations[n], locations[n + 1]))
        if (n + step) < order:
            G.add_edge(
                n, n + step,
                weight=distance(locations[n], locations[n + step]))

    print('size of graph G: {}'.format(G.size()))
    print(G.edges())
    return G

def distance(lhs, rhs):
    v = lhs - rhs
    return np.sqrt(np.dot(v, v))

def dump_spanning(T):
    print('size of MST: {}'.format(T.size()))
    for i in sorted(T.edges(data=True)):
        print(i)

if __name__ == '__main__':
    main()