Minimum spanning tree using Prim's Algorithm

mst

"""
  Produce a minimum spanning tree of the following graph. You do not need to show your work.
  (Source: A portion of the map of germany from the board game Power Grid.)
"""

def primsalgorithm(graph):
  if len(graph) == 0:
    return {}

  # Let v be an arbitrary vertex, let T be the tree consisting of just v
  v = graph.keys()[0]
  T = { v: [] }

  # While T is not a spanning tree
  while len(graph) != len(T):
    su = None
    sv = None
    sw = None

    # Look at all the edges in the cut induced by V(T)
    Tkeys = T.keys()
    for v in Tkeys:
      for e in graph[v]:
        # Let s = su-sv be the one with the smallest weight in the cut, say suEV(T), sv\EV(T)
        if (e[0] not in T) and ((sw == None) or (e[1] < sw)):
          su = v
          sv = e[0]
          sw = e[1]

    # Add su-sv to E(T), add sv to V(T)
    T[su].append((sv, sw)) # su is always EV(T)
    if T.get(sv):
      T[sv].append((su, sw))
    else:
      T[sv] = [(su, sw)]

  return T

germany = {
  'berlin': [('frankfurt', 6), ('magdeburg', 10), ('schwerin', 18), ('torgelow', 15)],
  'bremen': [('cuxhaven', 8), ('hamburg', 11), ('hannover', 10), ('osnabruck', 11), ('wilhelmshaven', 11)],
  'cuxhaven': [('bremen', 8), ('hamburg', 11)],
  'dortmund': [('essen', 4), ('kassel', 18), ('munster', 2)],
  'duisburg': [('essen', 0)],
  'dusseldorf': [('essen', 2)],
  'essen': [('dortmund', 4), ('duisburg', 0), ('dusseldorf', 2), ('munster', 6)],
  'flensburg': [('kiel', 4)],
  'frankfurt': [('berlin', 6)],
  'hamburg': [('bremen', 11), ('cuxhaven', 11), ('hannover', 17), ('kiel', 8), ('lubeck', 6), ('schwerin', 8)],
  'hannover': [('bremen', 10), ('hamburg', 17), ('kassel', 15), ('magdeburg', 15), ('osnabruck', 16), ('schwerin', 19)],
  'kassel': [('dortmund', 18), ('hannover', 15), ('osnabruck', 20)],
  'kiel': [('flensburg', 4), ('hamburg', 8), ('lubeck', 4)],
  'lubeck': [('hamburg', 6), ('kiel', 4), ('schwerin', 6)],
  'magdeburg': [('berlin', 10), ('hannover', 15), ('schwerin', 16)],
  'munster': [('dortmund', 2), ('essen', 6), ('osnabruck', 7)],
  'osnabruck': [('bremen', 11), ('hannover', 16), ('kassel', 20), ('munster', 7), ('wilhelmshaven', 14)],
  'rostock': [('schwerin', 6), ('torgelow', 19)],
  'schwerin': [('berlin', 18), ('hamburg', 8), ('hannover', 19), ('lubeck', 6), ('magdeburg', 16), ('rostock', 6), ('torgelow', 19)],
  'torgelow': [('berlin', 15), ('rostock', 19), ('schwerin', 19)],
  'wilhelmshaven': [('bremen', 11), ('osnabruck', 14)]
}

mst = primsalgorithm(germany)
print mst