pythonnumpy

Python/Numpy broadcast join between two arrays


The question is on how to join two arrays in this case more efficiently- There's a numpy array one of shape (N, M, 1) and array two of shape (M,F). It's required to join the second array with the first, to create an array of the shape (N, M, F+1). The elements of the second array will be broadcast along N.

One solution is copying array 2 to have size of the first (along all dims but one) and then concatenate. But this if the copying can be done as a broadcast during the join/concat it would use much lesser memory.

Any suggestions on how to make this more efficient?

The setup:

import numpy as np 

arr1 = np.random.randint(0,10,(5,10))
arr1 = np.expand_dims(arr1, axis=-1) #(5,10, 1)

arr2 = np.random.randint(0,4,(10,15)) 
arr2 = np.expand_dims(arr2, axis=0) #(1, 10, 15)

arr2_2 = arr2
for i in range(len(arr1)-1):
    arr2_2 = np.concatenate([arr2_2, arr2],axis=0)

arr2_2.shape #(5, 10, 15)
np.concatenate([arr1, arr2_2],axis=-1) # (5, 10, 16) -> correct end result

Joining arr1 and arr2 to get


Solution

  • try this

    >>> a = np.random.randint(0, 10, (5, 10))
    >>> b = np.random.randint(0, 4, (10, 15))
    >>> c = np.dstack((a[:, :, np.newaxis], np.broadcast_to(b, (a.shape[0], *b.shape))))
    >>> a.shape, b.shape, c.shape
    ((5, 10), (10, 15), (5, 10, 16)))