enumerate() array elements rather than indexing them in Python

by Forrest Sheng Bao http://fsbao.net

Recently, I am debugging two Python programs about graphs. I get very confused in many times, whether I am visiting an array element or the index of the array element. For example, I need to get the vertex number by indexing an array, and use that vertex number to index another array.

So, my friend Lay Yuan in Arizona State University told me a trick, using enumerate() in Python. This makes things much much more easier. I guess you are smart enough so this example can explain every thing to you

>>> for i,arr in enumerate([6,3,2,3,4]):
...     print i, arr
...
0 6
1 3
2 2
3 3
4 4

Below is an implementation of the Prim's algorithm. In this program, one step is to update the distance of current spanning tree to all neighbors of the newly added vertex. The neighbors of all vertexes are stored in an array neighbors. Without enumerate, I need these two lines to pick neighbors one by one:

for i in xrange(0,len(neighbors)):
w= neighbors[v][i];

With enumerate(), I just need one line to get the w:

for i, w in enumerate(neighbors[v]):

I don't need to struggle on the indexing problem any more.

# Python implementation of Prim's algorithm
# by Forrest Sheng Bao, Dept. of Computer Science, Texas Tech
# http://fsbao.net  forrest dot bao at gmail dot com
# This code implements the example on P. 571 of the famous
# "Introduction to Algorithms," 2nd ed., by Cormen, Leiserson, Rivest and Stein
# The connection matrix of the graph is represented as matrix A in the problem
# So it doesn't matter if you don't have the book. The vertexes are indexed from 0.

# The program is a free software. You may use, distribute and copy it under the 
# terms of GNU General Public License version 3.
# http://www.gnu.org/copyleft/gpl.html

# Expected output:
# $ python maxspanning.py
# vertex 0 has just been added into the growing spanning tree
# vertex 1 has just been added into the growing spanning tree
# vertex 2 has just been added into the growing spanning tree
# vertex 8 has just been added into the growing spanning tree
# vertex 5 has just been added into the growing spanning tree
# vertex 6 has just been added into the growing spanning tree
# vertex 7 has just been added into the growing spanning tree
# vertex 3 has just been added into the growing spanning tree
# vertex 4 has just been added into the growing spanning tree

from numpy import *

def spanningTree(n,A):
# neighbors[n]: // neighbors[i] is the list of neighbors of i in graph A
 neighbors=[];# neighbors[n], neighbors[i] is the list of neighbors of i in graph A
 for i in xrange(0,n):
  for j in xrange(0,n):
   neighbors.append([]);
   if( A[i,j] != 0):
    neighbors[i].append(j)

 treeNeighbor=[]; 
 # initially treeNeighbor[i] = -1;
 # during spanningTree(),
 #        treeNeighbor[i] == p  if selected[p] is true and  A[i,p]
 #            is the shortest distance from i to the current tree.
 #        treeNeighbor[i] == -1 if p is not a neighbor of any node of
 #            the current spanning tree.
 #        exception: 0 is always seleted first and treeNeighbor[0] is -1
 
 selected=[];#selected[i]=1: i is in spanning tree; 0: otherwise , initialized to be 0
 distance=[];#distance[i]: the shortest distance from i to tree, initialized to be +1
 
 for i in xrange(0,n):
  treeNeighbor.append(-1);
  selected.append(0);
  distance.append(100); # \infty for positive case, 1 for negative case

 first = 1; # a marker to mark, first =1 if the abitrary starting point is to be selected.
 for i in xrange(0,n):
  v = findMin(n,first,selected,treeNeighbor,distance);
  print "vertex", v, "has just been added into the growing spanning tree"

  first = 0;
  selected[v] = 1;    

#  add neighbors of v to treeNeighbor[] and update their distance to current spanning tree
  for i, w in enumerate(neighbors[v]):#each w \in neighbors[v] 
   #w= neighbors[v][i];
   if selected[w]==0 : #only update vertexes that are NOT on the tree now
    if A[v,w] < distance[w]:
     distance[w] = A[v,w];
     treeNeighbor[w] = v; #treeNeighbor, neighbors to current spanning tree
     #A[w,v]is the shortest path from w to current tree.
 return treeNeighbor;
 
def  findMin(n,first,selected,treeNeighbor,distance):
 if (first == 1):
  v = 0;
  treeNeighbor[0] = -1;
  return v ;
 else:
  min_distance=100; #\infty for postive case, 0 for negative case
  for v in xrange(0,n):
   if (selected[v]==0) and (treeNeighbor[v] != -1) and (distance[v] < min_distance):
   # the new node must be not on the current tree and be reachable to the tree in one hop and the shortest one to the current tree
    min_v = v;
    min_distance = distance[v]
  return min_v;

A=[\
[0, 4, 0, 0, 0, 0, 0, 8, 0],\
[4, 0, 8, 0, 0, 0, 0, 11, 0],\
[0, 8, 0, 7, 0, 4, 0, 0, 2],\
[0, 0, 7, 0, 9, 14, 0, 0, 0], \
[0, 0, 0, 9, 0, 10, 0, 0, 0],\
[0, 0, 4, 14, 10, 0, 2, 0, 0],\
[0, 0, 0, 0, 0, 2, 0, 1, 6],\
[8, 11, 0, 0, 0, 0, 1, 0, 7],\
[0, 0, 2, 0, 0, 0, 6, 7, 0]\
]

A=array(A)
R= spanningTree(9, A)