Enumerate

18 Oct 2010

This is a built-in Python function, but it's a very handy function that I didn't know about for some time, whilst wishing there was something exactly like it.

Python is great when it come to traversing lists with for loops, but sometimes you want to know where in the list you are. I used to write code something like this:

for n in range(len(my_list)):
    print n, my_list[n]

But it's much cleaner (and computationally more efficient, I believe) to use this:

for n, item in enumerate(my_list):
    print n, item

This function is particularly useful if you want to compare every item in a list to every other item in the list. For example, in my particle simulation, I wanted to test whether any particle in a list of my_particles overlapped with any other. The following code calls the collide function (which checks whether two particles overlap), with each pair of particles in the list.

for i, particle1 in enumerate(my_particles):
    for particle2 in my_particles[i+1:]:
        collide(particle1, particle2)

A further trick

You can also pass a second parameter, which defines what number to start counting from. For example:

a = ['two', 'three', 'four']
for i, word in enumerate(a, 2):
  print i, word

Will print:

2 two
3 three
4 four

Therefore the code for particle collisions above can be made a bit cleaner:

for i, particle1 in enumerate(my_particles, 1):
    for particle2 in my_particles[i:]:
        collide(particle1, particle2)

Python