Question

我有两个矩阵A和B，列数相等，但行数通常不等。我想将矩阵A和B中的所有列对相加。一个带有for循环的简单实现是这样的：

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

在这个例子中返回一个填充了1的50 * 3矩阵。感觉这应该可以通过一行代码实现。我怎样才能做到这一点？

Answer 1

它并不漂亮，但这条线路使用广播来做同样的事情：

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

实施例

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))

Answer 2

如果稍微重塑A，可以将两个矩阵一起添加并连接结果数组的第一个轴中的条目：

>>> 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.],
       ...

concatenate的替代方法稍微快一点就是使用reshape：

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

添加两个矩阵的所有列对

2 个答案: