Enumerate in Python

January 26, 2009

If you need to take an element of an array, do something to it, and store it in a new array you have some options. You could do this:

a = arange(10)
b = empty(10)
f = lambda x: x**2 #  or whatever
i = 0
for ai in a:
	b[i] = f(ai)
	i += 1

I think this is ugly because the i+=1 reminds me too much of Matlab. Maybe a bit better is:

a = arange(10)
b = empty(10)
f = lambda x: x**2
for i in range(len(a)):
	b[i] = f(a[i])

This is better, but I’m sad not to be iterating through a, and think the range(len(a)) construct is clumsy. We can have the best of both worlds using enumerate:

a = arange(10)
b = empty(10)
f = lambda x: x**2
for i,ai in enumerate(a):
	b[i] = f(ai)

Note that for these simple examples you could use b = [f(ai) for ai in a] (which works if a is any iterable object, like a list or a tuple) or simply b = f(a) (which works when a is a numpy array like above). Using enumerate often comes in handy for more complicated projects! For example, using the annoyingly indexed subplot command:

# X is a list of vectors
X = [rand(100) for i in range(10)]

for i,x in enumerate(X):
	# subplot is indexed using 1-indexing
	subplot(len(X),1,i+1) 
	plot(x)
Advertisements

Leave a Reply

Fill in your details below or click an icon to log in:

WordPress.com Logo

You are commenting using your WordPress.com account. Log Out / Change )

Twitter picture

You are commenting using your Twitter account. Log Out / Change )

Facebook photo

You are commenting using your Facebook account. Log Out / Change )

Google+ photo

You are commenting using your Google+ account. Log Out / Change )

Connecting to %s

%d bloggers like this: