Numpy quirk: Apply function to all pairs of two 1D arrays, to get one 2D array

Question

Let's say I have 2 one-dimensional (1D) numpy arrays, a and b, with lengths n1 and n2 respectively. I also have a function, F(x,y), that takes two values. Now I want to apply that function to each pair of values from my two 1D arrays, so the result would be a 2D numpy array with shape n1, n2. The i, j element of the two-dimensional array would be F(a[i], b[j]).

I haven't been able to find a way of doing this without a horrible amount of for-loops, and I'm sure there's a much simpler (and faster!) way of doing this in numpy.

Thanks in advance!

Answer 1

You can use numpy broadcasting to do calculation on the two arrays, turning a into a vertical 2D array using newaxis:

In [11]: a = np.array([1, 2, 3]) # n1 = 3     ...: b = np.array([4, 5]) # n2 = 2     ...: #if function is c(i, j) = a(i) + b(j)*2:     ...: c = a[:, None] + b*2  In [12]: c Out[12]:  array([[ 9, 11],        [10, 12],        [11, 13]]) 

To benchmark:

In [28]: a = arange(100)  In [29]: b = arange(222)  In [30]: timeit r = np.array([[f(i, j) for j in b] for i in a]) 10 loops, best of 3: 29.9 ms per loop  In [31]: timeit c = a[:, None] + b*2 10000 loops, best of 3: 71.6 us per loop