Adding up all pairs of columns of two matrices

0
votes

I have two matrices A and B with an equal amount of columns but generally an unequal amount of rows. I want to add up all pairs of columns from matrix A and B. A naive implementation with a for-loop is this:

import numpy as np
ncol = 3
nrow_A = 5
nrow_B = 10
A = np.ones((nrow_A,ncol))
B = np.zeros((nrow_B,ncol))

C = np.empty((nrow_A*nrow_B,ncol))
k = 0
for i in range(nrow_A):
    for j in range(nrow_B):
        C[k,:] = A[i,:]+B[j,:] 
        k += 1

which in this example returns a 50*3 matrix filled with ones. It feels like this should be possible with a single line of code. How can I achieve this?

2
pythonnumpymatrixcombinationsaddition

2 Answers

1
votes

If you reshape A slightly, you can add the two matrices together and concatenate the entries in the first axis of the resulting array:

>>> np.concatenate(A[:,np.newaxis,:] + B)
array([[ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       [ 1.,  1.,  1.],
       ...

A slightly quicker alternative to concatenate is to use reshape:

(A[:,np.newaxis,:] + B).reshape(-1, 3)
1
votes

It's not pretty but this line uses broadcasting to do the same thing:

(A[:, np.newaxis, :] + B[np.newaxis, :, :]).reshape(nrow_A * nrow_B, ncol)

Example

import numpy as np
ncol = 3
nrow_A = 5
nrow_B = 10
A = np.random.random_sample((nrow_A,ncol))
B = np.random.random_sample((nrow_B,ncol))

C = np.empty((nrow_A*nrow_B,ncol))
k = 0
for i in range(nrow_A):
    for j in range(nrow_B):
        C[k,:] = A[i,:]+B[j,:] 
        k += 1

D = (A[:, np.newaxis, :] + B[np.newaxis, :, :]).reshape(nrow_A * nrow_B, ncol)

print(np.allclose(C, D))