Minimum spanning trees tests

from collections import deque
import operator as op
import itertools

BIG = 1e10

class vertexColors:
    WHITE = 0
    GRAY = 1
    BLACK = 2

class vertex:
	color = vertexColors.WHITE
	distance = BIG
	pre = None
    
	def __init__(self,id):
		self.id = id
        
	def __repr__(self):
		return self.id
    
class edge:	
    def __init__(self,v,w,u=None):
        self.u = u
        self.v = v
        self.w = w
        
    def __repr__(self):
        return str(self.u) + "-(" + str(self.w) + ")-" + str(self.v)
    
    def __eq__(self, other):
		if other == None:
			return False
		if isinstance(other, edge):
			if self.u == other.u and self.v == other.v:
				if self.w == other.w:
					return True
				else:
					return False
			elif self.u == other.v and self.v == other.u:
				if self.w == other.w:
					return True
				else:
					return False
			else:
				return False
		else:
			return False
               
    
    def __hash__(self):
        """Very naive hashing, assumes vertex name is only a single letter"""
        u = ord(str(self.u))
        v = ord(str(self.v))
        
        if (u > v):
            h = u*30 + v
        else:
            h = v*30 + u
        return h
    
a = vertex("a")
b = vertex("b")
c = vertex("c")
d = vertex("d")
e = vertex("e")
f = vertex("f")
g = vertex("g")
h = vertex("h")
i = vertex("i")

wgraph = { a : [edge(b,4), edge(h,8)],
           b : [edge(a,4), edge(c,8), edge(h,11)],
           c : [edge(b,8), edge(d,7), edge(i,2), edge(f,4)],
           d : [edge(c,7), edge(e,9), edge(f,14)],
           e : [edge(d,9), edge(f,10)],
           f : [edge(c,4), edge(g,2), edge(d,14), edge(e,10)],
           g : [edge(f,2), edge(h,1), edge(i,6)],
           h : [edge(a,8), edge(b,11), edge(g,1), edge(i,7)],
           i : [edge(h,7), edge(g,6), edge(c,2)]
         }

def fix_edges(g):
    for v in g.keys():
        for e in g[v]:
            if e.u == None:
                e.u = v
        
def bfsw(g,s):
    
    for u in g.keys():
        u.color = vertexColors.WHITE
        u.distance = BIG
        u.pre = None

s.color = vertexColors.GRAY
    s.distance = 0
    s.pre = None
    q = deque()
    q.append(s)
    while len(q) > 0:
        u = q.popleft()
        for v,w in g[u]:
            if v.color != vertexColors.BLACK:
                if (u.distance + w) < v.distance:
					v.color = vertexColors.GRAY
					v.distance = u.distance + w
					v.pre = u
					q.append(v)
		u.color = vertexColors.BLACK

def make_edge_list(g):
    
    return set(list(itertools.chain(*g.values())))
       
def kruskal(g):
    A = set()
    sets = {}
    
    #create sets
    for v in g.keys():
        sets[v] = set([v])
    #sort edges
    edges = sorted(make_edge_list(g),key=op.attrgetter('w'))
	
    for edge in edges:
        if edge.v not in sets[edge.u] and edge.u not in sets[edge.v]:
            A.add(edge)
            u = sets[edge.v].union(sets[edge.u])
            for x in u:
                sets[x] = u
    return A

def prim(g,r):
    
    #set of all vertices
    G = set(g.keys())
	
    #set which will make tree
    T = set([r])
    #make cuts
    G.remove(r)
    
    #our MST
    E = set()
    
    #repeat while no more vertices are left
    guard = 1000
    while len(G) > 0 and guard > 0:
        w = BIG
        light_edge = None
        #choose light edge
    	for v in T:
            for e in g[v]:
                if e.w < w:
                    if e.u in T and e.v in T: #already in MST, ignore
                        continue
                    if e in E:    #already selected, ignore
                        continue
                    light_edge = e
                    w = e.w
                    
        #light edge chosen
        if light_edge == None:
            raise ValueError("couldnt find light edge")
        
        #print "selected ",light_edge
        
        T.add(light_edge.v)
        G.discard(light_edge.v)
        E.add(light_edge)
        guard = guard - 1
        
    if guard == 0:
        raise UserWarning("loop stuck")
        
    return E

fix_edges(wgraph)

#mst = kruskal(wgraph)
mst = prim(wgraph,a)
print mst

"""for v in wgraph.keys():
    print v.id+": "
    print "\t"+str(v.pre)
    print "\t"+str(v.color)
"""

Python Fiddle

Python Cloud IDE