CODEWITHSUNDEEP
pythonlistloops
DimKoim
DimKoim

Reputation: 1044

Pythonic way to write a loop

I have two lists: a = [1, 2, 3] and b = [4, 5, 6].

I have used two loops in python to subtract each element of b from each element of a.

import numpy as np
a = [1, 2, 3]
b = [4, 5, 6]
p = -1
result = np.zeros(len(a)*len(a))
for i in range(0,len(a)):
    for j in range(0,len(a)):
        p = p + 1
        result[p] = a[i] - b[j]

My result is correct: result = [-3., -4., -5., -2., -3., -4., -1., -2., -3.].

However, I would like to know if there is more elegant('pythonic') way to do it.

Upvotes: 2

Views: 17119

Answers (3)

Dima Tisnek
Dima Tisnek

Reputation: 11781

Since you are using numpy already, it's as simple as this:

In [29]: numpy.repeat(a, len(b)) - numpy.tile(b, len(a))
Out[29]: array([-3, -4, -5, -2, -3, -4, -1, -2, -3])

Per request in comments, OP wants 1d output size N^2 (although 2d might be more natural), for that numpy provides convenient functions to extend N-sized array to N^M, namely repeat and tile:

numpy.repeat([1,0], 2)
array([1, 1, 0, 0])

numpy.tile([1,0], 2)
array([1, 0, 1, 0])

Once both a and b are reformatted to the shape of p, it's a matter of element-wise subtraction that's native in numpy.

Upvotes: 5

Ashwini Chaudhary
Ashwini Chaudhary

Reputation: 250931

Pythonic way will be to use itertools.product, as it returns Cartesian product of the iterables passed to it. As no Python loop is involved it is going to be fast compared to version that are using loop:

>>> from operator import sub
>>> from itertools import starmap, product
>>> list(starmap(sub, product(a, b)))
[-3, -4, -5, -2, -3, -4, -1, -2, -3]

In NumPy you can do this using the recipes mentioned here:

>>> arr = np.dstack(np.meshgrid(a, b)).reshape(-1, 2)
>>> np.subtract(arr[:,0], arr[:,1])
array([-3, -2, -1, -4, -3, -2, -5, -4, -3])

Timing comparisons:

>>> b = [4, 5, 6]*1000
>>> a = [1, 2, 3]*1000
>>> %timeit list(starmap(sub, product(a, b)))
1 loops, best of 3: 464 ms per loop
>>> %timeit [x - y for x in a for y in b]
1 loops, best of 3: 491 ms per loop
>>> %%timeit                                 
result = []
for x in a:
    for y in b:
        result.append(x - y) #attribute lookup is slow
... 
1 loops, best of 3: 908 ms per loop
>>> %%timeit
result = [];append = result.append 
for x in a:
    for y in b:
        append(x - y)
... 
1 loops, best of 3: 617 ms per loop
#Numpy version will be little faster if a and b were nd arrays.
>>> %timeit arr = np.dstack(np.meshgrid(a, b)).reshape(-1, 2);np.subtract(arr[:,0], arr[:,1])
1 loops, best of 3: 573 ms per loop

Upvotes: 5

Matthias
Matthias

Reputation: 13222

There is no need to use an index. You can iterate over the values.

a = [1, 2, 3]
b = [4, 5, 6]
result = []
for x in a:
    for y in b:
        result.append(x - y)

The pythonic way would be a list comprehension.

a = [1, 2, 3]
b = [4, 5, 6]
result = [x - y for x in a for y in b]

Please bear in mind that you should use meaningful names for a, b, xand y in real code.

Upvotes: 14

Related Questions