CODEWITHSUNDEEP
python
will
will

Reputation: 10650

What is the pythonic way to loop through two arrays at the same time?

If I have two arrays, of the same length - say a and b

a = [4,6,2,6,7,3,6,7,2,5]

b = [6,4,6,3,2,7,8,5,3,5]

normally, I would do this like so:

for i in range(len(a)):
    print a[i] + b[i]

rather than something like this:

i=0
for number in a:
    print number + b[i]
    i += 1

because I prefer to be consistent with methods used.

I know of zip, but I never use it. Is this what zip was made for?

would 

for pair in zip(a,b):
    print pair[0] + pair[1]

be the pythonic way to do this?

Upvotes: 48

Views: 70385

Answers (5)

Muhammad Ali Malekzadeh
Muhammad Ali Malekzadeh

Reputation: 184

You can utilize a while loop to iterate over both arrays simultaneously.

i,j = 0,0
while i < len(a) and j < len(b):
    print(a[i] + n[j])
    i += 1 
    j += 1

Note that this loop is repeated for the shortest number of arrays.

Upvotes: 2

Rick
Rick

Reputation: 45231

Offering this answer for completeness since numpy has been discussed in another answer, and it is often useful to pair values together from higher ranked arrays.

The accepted answer works great for any sequence/array of rank 1. However, if the sequence is of multiple levels (such as a numpy array of rank 2 or more, but also such as in a list of lists, or tuple of tuples), one needs to iterate through each rank. Below is an example with a 2D numpy array: 

import numpy as np
a = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
b = np.array([list('abc'), list('pdq'), list('xyz')])
c = np.array([[frobnicate(aval, bval) for aval, bval in zip(arow, brow)] for arow, brow in zip(a, b)])

And the same concept will work for any set of two dimensional nested sequences of the same shape:

a = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
b = [list('abc'), list('pdq'), list('xyz')]
c = [[frobnicate(aval, bval) for aval, bval in zip(arow, brow)] for arow, brow in zip(a, b)]

If one or both of the nested sequences has "holes" in it, use itertools.zip_longest to fill in the holes (the fill value defaults to None but can be specified):

from itertools import zip_longest as zipl
a = [[], [4, 5, 6], [7, 8, 9]] # empty list in the first row
b = [list('abc'), list('pdq'), []] # empty list in the last row
c = [[frobnicate(aval, bval) for aval, bval in zipl(arow, brow)] for arow, brow in zipl(a, b)]

Upvotes: 2

fraxel
fraxel

Reputation: 35269

Instead of using zip you could use Numpy, especially if speed is important and you have long arrays. Its much faster and once you're using numpy arrays you don't need a loop, and can just write:

print a + b

Graph showing averaged timings for summing different length lists using zip, izip, and numpy: enter image description here

Upvotes: 9

unutbu
unutbu

Reputation: 879321

If the lists a and b are short, use zip (as @Vincenzo Pii showed):

for x, y in zip(a, b):
    print(x + y)

If the lists a and b are long, then use itertools.izip to save memory:

import itertools as IT
for x, y in IT.izip(a, b):
    print(x + y)

zip creates a list of tuples. This can be burdensome (memory-wise) if a and b are large.

itertools.izip returns an iterator. The iterator does not generate the complete list of tuples; it only yields each item as it is requested by the for-loop. Thus it can save you some memory.

In Python2 calling zip(a,b) on short lists is quicker than using itertools.izip(a,b). But in Python3 note that zip returns an iterator by default (i.e. it is equivalent to itertools.izip in Python2).

Other variants of interest:

Upvotes: 62

Vincenzo Pii
Vincenzo Pii

Reputation: 19805

A possible solution is using zip, as you mentioned yourself, but slightly differently than how you wrote it in the question:

for x, y in zip(a, b):
    print x, y

Notice that the length of the list of tuples returned by zip() will be equal to the minimum between the lengths of a and b. This impacts when a and b are not of the same length.

Upvotes: 13

Related Questions