More efficient of way of operating on numpy arrays

February 28, 2022

I have two numpy arrays, one with an array of 2 dimensional vectors (z1_i,z2_i)‘s and the other is just an array of 1 dimensional variables c_i‘s (the i’s here denote the index of the array). I want to operate on the two arrays to get a single array of one dimensional variable of the form z1_i*c_i + z2_i*(1-c_i) (again, the subscript i refer to the index of the array). For example:

2D_samples= np.array([[1, 2], [3, 4]])
1D_samples = np.array([1, 0])
# desired output would be np.array([1, 4])
# Specifically, [1*1+2*(1-1), 3*0+4(1-0)]=[1,4].

Is there a more efficient way of doing this without executing a for loop? Thanks.

Doing it by for loop would give:

output = np.zeros(2D_samples)
for i in range(len(2D_samples)):
  output[i] = 2D_samples[i][0] * 1D_samples[i] + 2D_samples[i][1] * (1 - c[i])

>Solution :

IIUC, you can do:

a2D = np.array([[1, 2], [3, 4]])
a1D = np.array([1, 0])

# reshape the 1D array to 2D
b = a1D[:,None]
# concatenate b and 1-b column-wise, multiply by a2D and sum per row
out = (a2D * np.c_[b,1-b]).sum(1)

output: array([1, 4])