#!/usr/bin/python
#
# COMP 4030 Project
# Chris Lawrence
# December 2, 1997
#
# Minimum Spanning Tree problem using Kruskal's algorithm for finding the
# MST of a graph.
#
# Graphs are internally represented in the adjacency list format.

from copy import copy
from whrandom import randint

# Two-dimensional array emulation
# From "Internet Programming with Python", Watters et al, p. 115
def multi_list(shape, initial_value = 0):
    if shape == []:
        return copy(initial_value)
    else:
        this_dimension = shape[0]
        others = shape[1:]
        array = [None]*this_dimension
        for i in xrange(this_dimension):
            array[i] = multi_list(others, initial_value)
        return array

# Disjoint set class - using path compression
class DisjointSet:
    def __init__(self, contents=[]):
        self.set = copy(contents);
        self.rank = len(contents)*[1]

    # Find out which set an item is in - find3 from text, p 179
    def find(self, x):
        r = x
        while self.set[r] != r: r = self.set[r]
        i = x
        while i != r:
            # Tuple-assignment saves a temporary variable
            (self.set[i], i) = (r, self.set[i])
        return r

    # Merge the sets - merge3 from text, p. 178
    def merge(self, a, b):
        if self.rank[a] == self.rank[b]:
            self.rank[a] = self.rank[a] + 1
            self.set[b] = a
        elif self.rank[a] > self.rank[b]:
            self.set[b] = a
        else:
            self.set[a] = b

# End of class DisjointSet

class ListGraph:
    def __init__(self, nodes = 0, edges = []):
        self.nodes = nodes
        self.edges = edges

    def __str__(self):
        return str(self.edges)

    def __repr__(self):
        return repr(self.edges)

    def __len__(self):
        return len(self.edges)

class AdjacencyList:
    def __init__(self, nodes):
        self.nodes = nodes
        self.edges = multi_list( [nodes, 0], [] )
        self.edgecount = 0

    def add_edge(self, a, b, weight):
        self.edges[a].append( (b, weight) )
        self.edges[b].append( (a, weight) )
        self.edgecount = self.edgecount + 1

    # Converts the adjacency list representation to a flat representation
    # which can be sorted
    def flatten(self):
        flat = []
        for i in range(self.nodes):
            for j in range(len(self.edges[i])):
                edge = self.edges[i][j]
                if i < edge[0]:
                    flat.append( (i, edge[0], edge[1]) )
        return ListGraph(self.nodes, flat)

    # String conversion of adjacency list
    def __str__(self):
        output = ""
        width = len(str(self.nodes-1))
        for i in xrange(self.nodes):
            output = output + ("%*d: %s\n" % (width, i, str(self.edges[i])))
        return output[:-1]

    def __repr__(self):
        return repr(self.edges)

    def __len__(self):
        return self.edgecount

# Compare the weights for edges
def compare(x, y):
    return cmp( x[2], y[2] )        

# Accepts a graph in adjacency-list format
# Output is in adjacency-list format
def Kruskal(g):
    graph = g.flatten() # Put into sortable array format
    # Sort using our comparison function
    # Python uses POSIX qsort() to sort lists
    graph.edges.sort(compare)
    n = graph.nodes
    T = AdjacencyList(n)
    lenT = 0 # Faster than doing len(T) each loop
    
    setstuff = n*[None] # Allocate space for n entries in the list
    for i in range(n):
        setstuff[i] = i;

    # Instantiate the DisjointSet class
    set = DisjointSet(setstuff)
    i = 0
    while lenT != n-1:
        try:
            (u, v, weight) = graph.edges[i]
        except IndexError:
            raise IndexError, "Disconnected graph"
            
        i = i + 1
        ucomp = set.find(u)
        vcomp = set.find(v)
        if ucomp != vcomp:
            set.merge(ucomp, vcomp)
            T.add_edge( u, v, weight )
            lenT = lenT + 1

    return T

# Run the algorithm for a graph with vertices in (u,v,w) format
def RunKruskal(g):
    print "Running Kruskal algorithm on the graph:"
    print g
    al = AdjacencyList(g.nodes)
    for edge in g.edges:
        al.add_edge( edge[0], edge[1], edge[2] )
    try:
        print "Result:\n", Kruskal(al)
    except IndexError:
        print "Graph is not connected, so no MST exists."

# If this code is executed on its own, run the test procedure
if __name__ == '__main__':
    g = ListGraph()
    g.nodes = 4
    g.edges = [ (0,2,2), (0,1,4), (0,3,9), (1,2,36), (2,3,3), (1,3,2) ]

    RunKruskal(g)
    print

    g = ListGraph()
    g.nodes = 4
    g.edges = [ (0,2,1), (0,1,48), (0,3,38), (1,2,32), (2,3,32), (1,3,42) ]
    
    RunKruskal(g)
    print

    g = ListGraph()
    g.nodes = 4
    g.edges = [ (0,1,48), (2,3,38) ]
    
    RunKruskal(g)
    print

    # Deallocate g
    del g

    n = 50
    print "%d-node graph test" % n
    al = AdjacencyList(n)
    # Make a complete graph
    for i in xrange(n):
        for j in xrange(0,i-1):
            al.add_edge(i, j, randint(1, 1000))
    try:
        print "Result:\n", Kruskal(al)
    except IndexError:
        print "Graph is not connected, so no MST exists."